build_llamalike_transformer

build_llamalike_transformer#

penzai.experimental.v2.models.transformer.variants.llamalike_common.build_llamalike_transformer(config: LlamalikeTransformerConfig, 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