def build_gpt_neox_attention(
    name: str,
    init_base_rng: jax.Array | None,
    config: GPTNeoXTransformerConfig,
) -> pz.nn.Attention:
  """Builds an attention block from a configuration.

  Args:
    name: Name of the attention block.
    init_base_rng: Base RNG for initializing the parameters.
    config: The configuration of the model.

  Returns:
    An Attention block.
  """
  num_heads = config.num_attention_heads
  embedding_dim = config.embedding_dim
  projection_dim = config.projection_dim

  # As used in https://github.com/google-deepmind/gemma.
  # (This exact value is probably not important.)
  masked_out_value = jnp.array(-2.3819763e38, dtype=config.activation_dtype)

  return pz.nn.Attention(
      input_to_query=pz.nn.Sequential([
          pz.nn.Affine.from_config(
              name=f"{name}/query",
              init_base_rng=init_base_rng,
              input_axes={"embedding": embedding_dim},
              output_axes={"heads": num_heads, "projection": projection_dim},
              dtype=config.parameter_dtype,
          ),
          pz.nn.ApplyRoPEToSubset(
              positions_input_name="token_positions",
              embedding_axis="projection",
              max_wavelength=config.rope_wavelength,
              rope_subset_size=config.rope_subset_size,
          ),
          pz.nn.ConstantRescale(
              by=jnp.array(projection_dim**-0.5, dtype=config.activation_dtype)
          ),
      ]),
      input_to_key=pz.nn.Sequential([
          pz.nn.Affine.from_config(
              name=f"{name}/key",
              init_base_rng=init_base_rng,
              input_axes={"embedding": embedding_dim},
              output_axes={"heads": num_heads, "projection": projection_dim},
              dtype=config.parameter_dtype,
          ),
          pz.nn.ApplyRoPEToSubset(
              positions_input_name="token_positions",
              embedding_axis="projection",
              max_wavelength=config.rope_wavelength,
              rope_subset_size=config.rope_subset_size,
          ),
      ]),
      input_to_value=pz.nn.Sequential([
          pz.nn.Affine.from_config(
              name=f"{name}/value",
              init_base_rng=init_base_rng,
              input_axes={"embedding": embedding_dim},
              output_axes={"heads": num_heads, "projection": projection_dim},
              dtype=config.parameter_dtype,
          ),
      ]),
      query_key_to_attn=pz.nn.Sequential([
          pz.nn.NamedEinsum(
              (
                  {"seq": "tq", "heads": "h", "projection": "p"},
                  {"seq": "tkv", "heads": "h", "projection": "p"},
              ),
              {"seq": "tq", "heads": "h", "kv_seq": "tkv"},
          ),
          pz.nn.ApplyCausalAttentionMask(
              masked_out_value=masked_out_value,
          ),
          pz.nn.Softmax("kv_seq"),
      ]),
      attn_value_to_output=pz.nn.Sequential([
          pz.nn.NamedEinsum(
              (
                  {"seq": "tq", "heads": "h", "kv_seq": "tkv"},
                  {"seq": "tkv", "heads": "h", "projection": "p"},
              ),
              {"seq": "tq", "heads": "h", "projection": "p"},
          ),
          pz.nn.Affine.from_config(
              name=f"{name}/output",
              init_base_rng=init_base_rng,
              input_axes={"heads": num_heads, "projection": projection_dim},
              output_axes={"embedding": embedding_dim},
              dtype=config.parameter_dtype,
          ),
      ]),
  )
Builds an attention block from a configuration.

  Args:
    name: Name of the attention block.
    init_base_rng: Base RNG for initializing the parameters.
    config: The configuration of the model.

  Returns:
    An Attention block.
  
penzai/models/transformer/variants/gpt_neox.py