Embeds shared state requests into a tree. This function is the inverse of hoist_shared_state_requests, and can be used to re-embed the initial va

Returns a stochastic, enabled version of this layer.

Returns a stochastic, enabled version of this layer.

Compatibility function: No-op for consistency with old Treescope API. Originally, global overrides for Penzai/Treescope ContextualValue objects n

Collects the effect types of all `EffectRequest` nodes in a (sub)model. Args: model_tree: A pytree for a model or submodel. Returns: A l

Embeds the given states into a handled model, and removes the handler. This function converts a functionalized stateful model into a single effectf

r"""Version of `freeze_variables` that only freezes `Parameter`\ s.

r"""Version of `freeze_variables` that only freezes `StateVariable`\ s.

Constructs a RMS layer normalization layer. Args: name: The name of the layer. init_base_rng: The base RNG to use for initializing mo

Constructs an ``EmbeddingTable``. Args: name: The name of the layer. init_base_rng: The base RNG to use for initializing model parame

Constructs an ``AddBias`` layer from a configuration. Args: name: The name of the layer. init_base_rng: The base RNG to use for initi

Constructs an uninitialized Gemma feed-forward layer. Gemma's feedforward layer uses GELU-based gated linear units (GEGLU), as proposed by Sh

Builds a ``GemmaTransformerBlock`` from a configuration. Args: config: The configuration of the Gemma model. Returns: A ``GemmaT

Constructs an uninitialized transformer with the Gemma architecture. Args: config: The configuration of the Gemma model. Returns:

Builds a TellIntermediate layer that writes to the given tag.

Constructs a RMS layer normalization layer. Args: across_axes: Names and lengths of the axes to normalize over. epsilon: Epsilon para

Creates an ``ApplyAttentionMask`` layer from a tag and a mask value. Args: mask_tag: Side input tag for the mask side input. This should be

Constructs an ``EmbeddingTable`` with uninitialized parameters. Args: vocab_size: The size of the vocabulary. embedding_axes: A dicti

Constructs an ``AddBias`` layer from a configuration. Args: biased_axes: Names and lengths for the axes in the input that the bias

Constructs a ``GemmaTransformer`` from an official Gemma Flax checkpoint. The parameters of the loaded ``GemmaTransformer`` will be close to thos

Transforms a `Transformer` into cached sampling mode. This constructor hot-swaps all `pz.nn.Attention` layers in the original model to enable

Transforms a `GemmaTransformer` into cached sampling mode. This constructor hot-swaps all `model_core.GemmaAttention` layers in the original

Builds a caching attention from an uncached attention. Args: original: The original attention layer that this block should replace. s

Builds a ``ShareableUninitializedParameter`` equivalent to ``uninit``.

Constructs a full named array with a given shape. Args: named_shape: Named shape for the result. fill_value: Value to fill the array with.

Substitutes all dimension variables for their known values. Args: structure: A PyTree structure containing `ArraySpec` leaves. substitution

Builds a Gemma model from a pretrained checkpoint. The parameters of the loaded ``Transformer`` will be close to those in the original checkpoint

Returns a collection of all dimension variables in a structure. Args: structure: The structure to extract variables from. Usually a PyTree with

Returns the collection of selected key paths.

Returns the slot that this variable is replaced with when unbound. The variable slot will uniquely identify this variable within a pytree.

Returns the slot that this variable is replaced with when unbound. The `get_slot` method of a frozen variable should always return the same s

Converts a GPT-NeoX model to a Penzai model. This function converts GPT-NeoX models from their HuggingFace implementations to Penzai. Note: Ch

Renders ArraySpec and contained variables.

Handles data effects objects.

Renders a penzai layer. Layers render like Structs in general, except that they render with input and output structure annotations. Args:

Renders a penzai layer. Layers render like Structs in general. However: * The outermost layer is augmented with information about the set of fr

Renders NamedArrays.

Renders a penzai struct or layer. Differences in handling structs vs other dataclasses: * We use ``cls.from_attributes`` instead of ``penzai

Returns the ID of the handler that created this implementation.

Builds a ``CollectingSideOutputs`` that handles effects in this layer. Args: body: The layer to wrap. Usually will contain `SideOutputReque

Builds a `WithSideInputsFromInputTuple` handling effects in this layer. Args: body: The layer to wrap. Usually will contain `SideInputReque

Builds a ``WithConstantSideInputs`` that handles effects in this layer. Args: body: The layer to wrap. Usually will contain `SideInputReque

Builds a `WithStatefulRandomKey` that handles effects in this layer. Args: body: The layer to wrap. Usually will contain random effects in

Builds a `WithFrozenRandomState` that handles effects in this layer. Args: body: The layer to wrap. Usually will contain random effects in

Extracts all constant side inputs from a tree so they can be re-handled. This function finds all `WithConstantSideInputs` handlers in a tree, and

r"""Hoists out the value of shared states in a pytree. This function is a helper for manipulating Penzai models that contain shared states. Ordin

Tries to generate a unique handler ID from the structure of a subtree. Each handler in Penzai's effect system needs to have a unique ID that iden

Initializes all parameters in a model. Args: model: Model to initialize. prng_key: Key to use to initialize parameters. Returns: Cop

Initializes the parameters of a model, sharded according to a mesh. Args: model: A model whose parameters we should initialize. Should usually

Inlines instances of `Sequential` (not subclasses) into parent groups.

Inlines instances of `Sequential` (not subclasses) into parent groups.

The axis names and sizes that should appear in the input only.

The axis names and sizes that should appear in the input only.