build_llamalike_feedforward

build_llamalike_feedforward#

penzai.models.transformer.variants.llamalike_common.build_llamalike_feedforward(name: str, init_base_rng: jax.Array | None, config: LlamalikeTransformerConfig) → model_parts.TransformerFeedForward[source]#

Creates a feedforward block.

This family of models use gated linear units, as proposed by Shazeer (2020). We represent this computation as a composition of simpler Penzai primitives, to enable patching and post-processing of the various internal activations.

Parameters:

  • name – Name of the feedforward block.

  • init_base_rng – Base RNG for initializing the parameters.

  • config – The configuration of the model.

Returns:

An instance of TransformerFeedForward containing the GELU MLP blocks.

Contents