def build_gpt_neox_feedforward(
    name: str,
    init_base_rng: jax.Array | None,
    config: GPTNeoXTransformerConfig,
) -> model_parts.TransformerFeedForward:
  """Creates a feedforward block.

  The GPT-NeoX model uses a standard MLP configuration.

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

  Returns:
    An instance of TransformerFeedForward containing the MLP.
  """
  act_fn = {
      "gelu_exact": functools.partial(jax.nn.gelu, approximate=False),
      "gelu_approx": functools.partial(jax.nn.gelu, approximate=True),
      "silu": jax.nn.silu,
      "relu": jax.nn.relu,
  }[config.activation_fn]
  return model_parts.TransformerFeedForward([
      pz.nn.Affine.from_config(
          name=f"{name}/in",
          init_base_rng=init_base_rng,
          input_axes={"embedding": config.embedding_dim},
          output_axes={"neurons": config.mlp_hidden_dim},
      ),
      pz.nn.Elementwise(act_fn),
      pz.nn.Affine.from_config(
          name=f"{name}/out",
          init_base_rng=init_base_rng,
          input_axes={"neurons": config.mlp_hidden_dim},
          output_axes={"embedding": config.embedding_dim},
      ),
  ])
Creates a feedforward block.

  The GPT-NeoX model uses a standard MLP configuration.

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

  Returns:
    An instance of TransformerFeedForward containing the MLP.
  
penzai/models/transformer/variants/gpt_neox.py