build_gpt_neox_feedforward

build_gpt_neox_feedforward#

penzai.experimental.v2.models.transformer.variants.gpt_neox.build_gpt_neox_feedforward(name: str, init_base_rng: jax.Array | None, config: GPTNeoXTransformerConfig) → model_parts.TransformerFeedForward[source]#

Creates a feedforward block.

The GPT-NeoX model uses a standard MLP configuration.

Parameters:

  • 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.

Contents