A leaf type for use in tests.

(Visually-distinctive) marker class for the selected subtree in tests.

A simple identity layer with a comment tag, for assertions.

A simple identity layer with a comment tag, for assertions.

r"""A multidimensional array with a combination of positional and named axes. Conceptually, ``NamedArray``\ s can have positional axes like an ordi

Helper object for building a selection. Not intended to be exposed to user code! If you're seeing this and you don't expect to, please file an is

Reference for a handled random effect.

Helper object to call a function with it's value and a child. Attributes: child: A PyTree child. Can be traversed to handle effects. fn: A

A dimension variable, which will be matched with the shapes in a structure. In a slight abuse of syntax, ``DimVar`` can be unpacked using ``*`` or

A possibly-transposed view of an array with positional and named axes. This view identifies a particular set of axes in a data array as "virtual

A marker object used to temporarily assign names to positional axes. Every ``TmpPosAxisMarker`` is unique, comparing equal only to itself, so it is

A constraint that we have already failed to satisfy.

A selected subset of nodes within a larger PyTree. Penzai selectors (such as ``.at(...)``) return ``Selection`` objects, which indicate a specifi

Marker for a handled local state effect. Attributes: handler_id: The ID of the handler that is responsible for handling this effect.

A dimension with a known value, used to bind a name to a value directly. The usual way to create a ``KnownDim`` is by calling `vars_for_axes` with

A learnable parameter. Most initialized layers should use Parameter to store their learnable parameters. Attributes: value: The value of t

Error raised due to pytree dataclass safety checks.

Implementation of the random number generation effect.

Raised when a `check_structure` call fails.

A binding for a variable set by a keypath.

A non-leaf marker for a (named) array structure. This is like a `jax.ShapeDtypeStruct`, but it is an empty PyTree node instead of being a leaf, s

Marker for a global attention block.

Wrapper for a module that automatically applies `nmap` to callables. Attributes: _module: The module to wrap.

A no-op layer taking the place of a disabled `StochasticDropout` layer. This layer can be used to mark a location in a model where dropout could

A no-op layer taking the place of a disabled `StochasticDropout` layer. This layer can be used to mark a location in a model where dropout could

Effect request for local state with a frozen value. This can be used to store the state for a stateful model after it has been initialized and po

Sampling state for the key-value-caching Gemma variant. You should not usually need to construct this on your own. Instead, it will be returned b

Reference for a handled side input effect.

A frozen representation of data in a particular Flax scope. Flax implements its modules using a "functional core" which is a stateful manager of

A frozen representation of data in a particular Flax scope. Flax implements its modules using a "functional core" which is a stateful manager of

Internal state for the trainer. Attributes: step: Current step of training. opt_state: An optimizer state for the learnable parts of `model

A variable standing in for a sequence or dictionary axis names. The usual way to create a ``MultiDimVar`` is by unpacking the `var` function, e.g

A slot for a parameter in a model. A ParameterSlot identifies a place in a model (or other pytree) where a Penzai parameter should be substituted

Sentinel marker for the remainder of axes.

State that manages the set decoded tokens during sampling. The purpose of this class is to keep the outputs, inputs, and key-value caches of the

A layout mark that can be used to identify selected nodes.

Effect request for local state that is shared. A ``SharedLocalStateRequest`` can be used to share an explicitly-named state variable between mult

Implementation of the side input effect.

Effect request for a side output. Since side outputs do not change the behavior of the model, ``SideOutputRequest`` implements `tell` (but in a w

A slot for a parameter in a model. A StateVariableSlot identifies a place in a model (or other pytree) where a Penzai state variable should be su

Stochastic dropout layer. Dropout layers randomly mask out elements with a probability ``drop_rate``, and then scale the output up by a factor of

Stochastic dropout layer. Dropout layers randomly mask out elements with a probability ``drop_rate``, and then scale the output up by a factor of

Container for a struct's static fields. Attributes: child_field_names: Names of all non-static fields in PyTree serialization order.

Raised when attempting to access the value of an unbound variable.

Exception raised when a method is called on an unhandled effect.

Raised when a Variable label is used by multiple Variables.

A constraint for matching two named shapes.

A constraint for matching two positional shapes.

Adjusts input embeddings using rotary position embeddings (RoPE). Rotary position embeddings, proposed by Su et al. (2021), incorporate relative

Representation of the arguments and kwargs of a function call.

Representation of the arguments and kwargs of a function call.

Marker for a local sliding-window attention block. Attributes: window_size: Size of the sliding window.

A label for a StateVariable that is unique based on Python object ID.

No-op annotation that indicates values of many different shapes.

A substitution for all of the dimension variables in a structure. DimensionVariableSubstitution can be used to inspect and modify the dimension v

A StateVariable label for a Flax variable.

Configuration parameters for a GPT Neo-X transformer. These are held in a single configuration object to simplify argument passing during constru

Input structure for `GemmaTransformer`. Attributes: tokens: Sequence of tokens, as an integer named array with a "seq" axis and possibly

Input structure for the `GemmaKVCachingTransformer`. Attributes: tokens: Subsequence of the current tokens we are processing, as an integer

Top-level Gemma transformer decoder, encapsulating all internal effects. This class represents the full Gemma model, and can be loaded from the o

Common configuration parameters for the Gemma transformer architecture. These are held in a single configuration object to simplify argument passin

Marker for a handled local state effect. Attributes: handler_id: The ID of the handler that is responsible for handling this effect.

No-op annotation that holds onto intermediate activations.

A representation of an intercepted Flax module method call. An InterceptedFlaxModule captures the logic that runs when you call a single Flax mod

A representation of an intercepted Flax module method call. An InterceptedFlaxModule captures the logic that runs when you call a single Flax mod

An isolated part of a model, with saved inputs, outputs, and variables. Variable values will also be frozen at the state they had when the model

An isolated part of a submodel, with its saved inputs and outputs. When initially populated by `isolate_submodel`, any states used by the submode

Helper object that captures information while running the submodel.

Helper object that captures information necessary to isolate the submodel.