build_gpt_neox_transformer

build_gpt_neox_transformer#

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

Builds a Llama-like transformer model from a configuration.

Parameters:

  • config – The configuration of the model.

  • init_base_rng – Base RNG for initializing the parameters.

  • name – Name for the top-level model, used as a prefix for all parameters.

Returns:

A full transformer model.

Contents