Helper to partition the indices of a dict-style index.

Inserts `_InterceptShapes` layers around all sublayers of a layer.

Inserts inserted_cls layers around a layer and inside sublayers.

Applies a function, passing both the selected nodes and their indices. Indices are taken relative to the linearized sequence of currently selecte

Selects nodes by their keypaths (relative to the current selection). Args: keypaths: A collection of keypaths. Returns: A new se

Broadcasts a named array to be compatible with another. Args: other: Another named array. Returns: A named array that has the sa

Checks that indices appear exactly once in the given axis specs.

Checks that indices appear exactly once in the axis specifications.

Returns the protocol that this effect request handles.

Returns an `ArraySpec` with this named shape and `np.floating` dtype.

Replaces each variable in a pytree with a frozen copy. The resulting tree will contain frozen variable instances instead of mutable variable inst

Directly instantiates a struct given all of its fields. Structs can override ``__init__`` to have arbitrary custom behavior, but this may mak

Constructs a ``Linear`` layer from a configuration. This can be used when building a new linear operator at the start of training. Note:

Constructs an ``ApplyRoPE`` layer for a given axis and side input tag. Args: positions_tag: Side input tag for the position side input. Thi

Constructs a ``Linear`` layer from a configuration. This can be used when building a new linear operator at the start of training. The return

Get values from the array.

Constructs the initial training state. Args: model: The model being optimized. optimizer_def: The optax optimizer to use. root_

Returns True if this field is treated as a PyTree child node by Struct. Fields are treated as PyTree nodes by default unless they contain a metadat

Checks if a tree is a subclass of `Sequential` or a `NamedGroup`.

Checks if a tree is a subclass of `Sequential` or a `NamedGroup`.

Summarizes a (validly constructed) named array.

Gets the next key from this stream, mutating the stream in place.

Partitions the tree into ``(selected_tree, remainder_tree)`` parts. This function can be used to separate out the selected components of a tree

Filters a selection to only the nth selected node. ``n`` is taken relative to the linearized sequence of currently selected nodes, in the sen

Runs one step of training.

Computes a CSS color to display for this object in treescope. This function can be overridden to change the color for a particular object in

Activates the random stream, returning itself for convenience.

Converts a view into a proper `NamedArray` by moving positional axes. The resulting `NamedArray` has the same named and positional shapes as this

Identifies whether a token is a control token.

Mutable variable protocol method, for JAX ecosystem interoperability.

Mutable variable protocol method, for JAX ecosystem interoperability.

Mutable variable protocol method, for JAX ecosystem interoperability.

Renders this selection as the root object in a treescope rendering.

Helper function to get an `ArraySpec` view of a leaf.

Creates a summary line for an array structure.

Helper to call a callable pytree with an argument.

Helper to call a callable pytree with an argument.

Checks if we need special handling for this subtree.

JIT-compiled helper for getitem.

JIT-compiled helper for in-place updates.

Helper function to match a leaf value against an `ArraySpec`. Args: array_structure: Expected structure to match agains. value: Value to ma

Simplifies a named constraint by inlining multivars.

Generates a unique label for a Variable.

Simplifies a positional constraint by inlining multivars.

Helper function to rename all `SupportsParameterRenaming` subtrees.

Inserts `_InterceptShapes` layers around a layer.

Inserts _InterceptValues layers around a layer.

Advances a sampling state by one token. This can be used to feed new sampled tokens one-at-a-time through the model, producing new log-probs that

Collects the set of all handler IDs inside a model or submodel. Args: model_tree: A pytree for a model or submodel. Returns: A list of a

Creates a placeholder parameter containing `jax.ShapeDtypeStruct` leaves. This can be used to create a model with the right PyTree structure with

Converts into a `NamedArrayView`, keeping positional axes. This function is usually not necessary for ordinary named-array manipulation, sinc

Selects all PyTree nodes with no children, including PyTree leaves. This is different than `at_pytree_leaves` in that it additionally selects

Constructs a dictionary with all of the fields in the class. The result of this should be passable back to `from_attributes` to rebuild (a co

Binds variables (mutable or frozen) into the variable slots in a pytree. This function re-inserts variable instances into a tree in place of their

Builds an attention block from a configuration. Args: name: Name of the attention block. init_base_rng: Base RNG for initializing the param

Builds a GPT-NeoX "parallel" transformer block from a configuration. GPT-NeoX uses a parallel formulation of transformer blocks, where the input of

Creates a feedforward block. The GPT-NeoX model uses a standard MLP configuration. Args: name: Name of the feedforward block. init_base_

Builds a Llama-like transformer model from a configuration. Args: config: The configuration of the model. init_base_rng: Base RNG for initi

Builds an attention block from a configuration. Args: name: Name of the attention block. init_base_rng: Base RNG for initializing the param

Builds a transformer block from a configuration. Args: name: Name of the block. init_base_rng: Base RNG for initializing the parameters.

Creates a feedforward block. This family of models use gated linear units, as proposed by Shazeer (2020). We represent this computation as a comp

Builds a Llama-like transformer model from a configuration. Args: config: The configuration of the model. init_base_rng: Base RNG for initi

Checks that there are no duplicated parameters in a model. Args: model: An arbitrary model. Raises: ValueError: If two concrete paramete

Returns an initializer that uses a constant value.

Returns an initializer that uses a constant value.

Contracts two named arrays along the given axis names. Args: names: The axis names to contract. Can also be a single string axis name. left

Contracts two named arrays along the given axis names. Args: names: The axis names to contract. Can also be a single string axis name. left

Derives a PRNG key for a parameter from a base key and a name. Args: base_rng: The base PRNG key. name: The name of the parameter. Retur

Returns the protocol that this effect request handles.

Returns the effect protocol(s) that this handler handles. Advanced handlers are allowed to handle multiple effects, and the specific effect i

Flattens a selection whose selected values are all selections. This function takes a selection for which all of the selected values are alrea

Returns a frozen copy of this variable.

Constructs a simple input structure for a batch of single segments. This can be used to process inputs that do not need advanced position or

Constructs a simple input structure for a batch of unpadded samples. This can be used to process inputs that do not need advanced position or

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

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

Builds a LoRA layer from a `pz.nn.Linear` layer. Args: linear: The linear layer to adapt. name: Name for this layer's parameters. Mus

Builds a LoRA layer from a `pz.nn.Linear` layer. Args: linear: The linear layer to adapt. rank: The rank of the low-rank adapter.

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