build_gpt_neox_attention

build_gpt_neox_attention#

penzai.experimental.v2.models.transformer.variants.gpt_neox.build_gpt_neox_attention(name: str, init_base_rng: jax.Array | None, config: GPTNeoXTransformerConfig) → pz.nn.Attention[source]#

Builds an attention block from a configuration.

Parameters:

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

Contents