def build_gpt_neox_block(
    name: str,
    init_base_rng: jax.Array | None,
    config: GPTNeoXTransformerConfig,
) -> model_parts.TransformerBlock:
  """Builds a GPT-NeoX "parallel" transformer block from a configuration.

  GPT-NeoX uses a parallel formulation of transformer blocks, where the input of
  the previous block is fed to the attention and feedforward components at the
  same time:

    y = x + MLP(LayerNorm(x)) + Attention(LayerNorm(x))

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

  Returns:
    A full transformer block.
  """
  return model_parts.TransformerBlock(
      sublayers=[
          pz.nn.BranchAndAddTogether([
              pz.nn.Identity(),
              pz.nn.Sequential([
                  pz.nn.LayerNorm.from_config(
                      name=f"{name}/pre_attention_norm",
                      init_base_rng=init_base_rng,
                      across_axes={"embedding": config.embedding_dim},
                      epsilon=config.layernorm_epsilon,
                      dtype=config.parameter_dtype,
                  ),
                  build_gpt_neox_attention(
                      f"{name}/attention", init_base_rng, config
                  ),
              ]),
              pz.nn.Sequential([
                  pz.nn.LayerNorm.from_config(
                      name=f"{name}/pre_ffw_norm",
                      init_base_rng=init_base_rng,
                      across_axes={"embedding": config.embedding_dim},
                      epsilon=config.layernorm_epsilon,
                      dtype=config.parameter_dtype,
                  ),
                  build_gpt_neox_feedforward(
                      f"{name}/mlp", init_base_rng, config
                  ),
              ]),
          ])
      ],
  )
Builds a GPT-NeoX "parallel" transformer block from a configuration.

  GPT-NeoX uses a parallel formulation of transformer blocks, where the input of
  the previous block is fed to the attention and feedforward components at the
  same time:

    y = x + MLP(LayerNorm(x)) + Attention(LayerNorm(x))

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

  Returns:
    A full transformer block.
  
penzai/models/transformer/variants/gpt_neox.py