build_llamalike_block

build_llamalike_block#

penzai.models.transformer.variants.llamalike_common.build_llamalike_block(name: str, init_base_rng: jax.Array | None, config: LlamalikeTransformerConfig, block_index: int | None = None) → model_parts.TransformerBlock[source]#

Builds a transformer block from a configuration.

Parameters:

  • name – Name of the 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:

A full transformer block.

Contents