build_llamalike_attention

build_llamalike_attention#

penzai.experimental.v2.models.transformer.variants.llamalike_common.build_llamalike_attention(name: str, init_base_rng: jax.Array | None, config: LlamalikeTransformerConfig, block_index: int | None = None) → 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.

  • block_index – The index of the transformer block in the list of blocks. Can be None if the attention type doesn’t depend on the block index.

Returns:

An Attention block.

Contents