Wraps a positional array as a ``NamedArray``. Args: array: Array to wrap. *names: Optional names for the axes of the array. If provid

Wraps this struct in a selection, enabling functional-style mutations. This is a convenience wrapper around selectors.select to enable easier

Selects subtrees that are an instance of the given type. Convenience wrapper for:: .at_subtrees_where(lambda subtree: isinstance(subtree,

Ensures that the named axes are stored in a canonical order. Returns: Equivalent `NamedArray` whose data array contains the positional axes

Wraps an array method on a named array.

Builds a redirector module for the given target module.

Builds a nmap-wrapped function with a docstring.

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

Returns the result of a singleton selection. Returns: The selected subtree from this selection. Raises: ValueError: If this sele

Attaches names to the first positional axes in an array or view. This is a version of `tag` that allows you to name only a subset of the arra

Updates the value of this variable to match a frozen variable.

Selects a specific child of each selected node. ``Selection.at(...)`` allows you to modify a tree with an almost-imperative style while maint

Gets the selected subtree(s) in order. Convenience wrapper for ``.selected_by_path.values()``. Returns: A tuple containing the selecte

Rebuilds the tree, forgetting which nodes were selected. Returns: A copy of `remainder` with the holes filled by the values in `selec

Swaps the order of operations for a binary operation.

Replaces the selected subtree(s) with a fixed replacement. Args: replacement: The pytree to replace with. Returns: A modified ve

Adds the requested axes to the front of the array's positional axes. This is a version of `untag` that can be called on NamedArrays or NamedA

Summarizes a dimension variable into a human-readable summary.

Sets the current state of the local variable.

Internal function to build a selection from a boundary. This function, like the `_InProgressSelectionBoundary` class, is an implementation detail

Returns the number of elements in the selection.

Builds a GemmaAttention block from a configuration. Args: config: The configuration of the Gemma model. Returns: A GemmaAttentio

Ensures that this array's PyTree structure matches another array's. This can be used to ensure that one named array has the same PyTree struc

Updates values in the array.

Selects all direct children of each selected subtree. This can be used to implement recursive tree traversals in a generic way, using somethi

Convenience function to bind variables to a layer. `layer.bind_variables(variables)` is a simple alias for `pz.bind_variables(layer, variable

Gets the current state of the local variable.

Returns the next key in the sequence, and advances the stream.

Builds a tuple that can be used to compare this object to other types.

Retrieves the value for the side input.

Summarizes a pattern into a human-readable summary.

Context manager to add informative details to selection errors.

Captures positional and keyword arguments from this call.

Returns a new random stream with the given base key.

Writes a value to the side output.

Internal helper to match a pattern with a value. Args: keypath: Keypath for the current value. pattern: Pattern to match. value: Value

Wraps a scalar conversion operator on a named array.

Selects subtrees that are equal to a particular object. Mostly a convenience wrapper for :: .at_subtrees_where(lambda subtree: template ==

Randomly initializes the parameter. Args: prng_key: Key to use for initialization. Returns: A new instance of ``Parameter``.

Directly initializes the parameter with particular value. This can be used to bypass the initializer function and set the value directly, whi

Returns True if the selection is empty.

Shards a tree of `pz.nx.NamedArray` objects based on their axis names. Args: tree: A PyTree of `pz.nx.NamedArray` instances, with the same stru

Automatically vectorizes ``fun`` over named axes of NamedArray inputs. ``nmap`` is a "named-axis vectorizing map". It wraps an ordinary positiona

Ensures that the named axes are stored in this order, keeping them named. This function can be used if it is important for the axis names to appe

Replaces the selected subtree(s) based on their path(s). If you need both the value and the key, see ``.apply(fn, with_keypath=True)``.

Calls a layer with temporary variables, without modifying its state. This is a convenience method for: * freezing any variables currently

Extends the default renderer to support __penzai_repr__.

Checks that an expected number of nodes is selected. Args: count: The expected number of nodes. Returns: The original selection

Selects all PyTree leaves of each selected subtree. This selects all of the leaves of the PyTree according to `jax.tree_util`, giving the mos

Builds a layer stack of layers with non-shared parameters. This function assumes that all variables returned by this builder are defined insi

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

Retrieves the selected subtree(s) based on their path(s). Returns: A dictionary of selected nodes, indexed by their path

Gets the default color for a given effect (for treescope rendering).

Returns the input structure of this layer. The input structure of a layer is a PyTree describing the structure the layer expects to be called

Inverts a selection, selecting subtrees with no selected children. ``selection.invert()`` selects the largest set of subtrees such that those

Flattens this tree node. See `jax.tree_util.register_pytree_with_keys_class`. This method should not be overridden by subclasses, since

Decodes a token ID to a string.

Inserts inserted_cls layers around sublayers of a layer.

Inserts _InterceptValues layers around sublayers of a layer.

Broadcasts a named array to a possibly-larger shape. Args: positional_shape: Desired positional shape for the array. Will be broadc

Checks that the names in the array are correct.

Returns the protocol that this effect request handles.

Converts an ``ArraySpec`` into a (possibly wrapped) PyTree leaf. By default, an ``ArraySpec`` has no PyTree children. This method can be used

Returns the output structure of this layer. The output structure of a layer is a PyTree describing the structure of the layer's return value,

Wraps a tree in a singleton selection for processing. Args: tree: A tree to select. Returns: A selection of the entire tree. Useful beca

Replaces the selected subtree(s) in order. Args: replacements: An iterable of PyTrees to insert at the selected locations, in order

Attaches names to the positional axes of an array or view. Args: *names: Axis names to assign to each positional axis in the array or view.

Helper function to preprocess token arrays and axis names.

Computes query, key, and value head axes and einsum names.