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

Unflattens this tree node.

Shape-checking opt-out decorator for ``Layer.__call__``.

Creates an `InterceptedFlaxModuleMethod` from applying a Flax module. Note that this function is intended for interactive exploration and to help

Creates an `InterceptedFlaxModuleMethod` from applying a Flax module. Note that this function is intended for interactive exploration and to help

Unpatches the IPython debugger so that it uses its original displayhook.

Splits a named array across a given named axis. Args: array: The array to split. axis_name: Name of the axis along which to unstack. Ret

Produces a positional view of the requested axis names. `untag` can only be called on a `NamedArray` or `NamedArrayView` that does not have a

Unwraps this array, possibly mapping axis names to positional axes. Unwrap can be called either on arrays with no named axes, or arrays with

Raises an error when accessed.

Raises an error when accessed.

The value of the parameter. May raise an error for some instances.

The value of the parameter. May raise an error for some instances.

r"""Value accessor for compatibility with `ParameterLike`. Raises: UninitializedParameterError: Since ``UninitializedParameter``\ s are not

Value accessor for compatibility with `ParameterLike`. Raises: MissingSharedParameterError: Since ``SharedParameterLookup`` does not have

The structure of the parameter.

Creates a variable for an axis shape. The resulting variable can be used in place of a single array dimension. It can also be unpacked with ``*va

Variable-aware version of `jax.jit`. This function is like `jax.jit`, but adds support for Variables as leaves of the input pytree(s). Limitat

Generic variance scaling initializer.

Generic variance scaling initializer.

Creates variables for a known collection of named axes. Args: var_name: A name for the variable that will store the concrete sizes for all

Ensures a view is a `NamedArray` by moving positional axes. The resulting `NamedArray` has the same named and positional shapes as this objec

Returns a copy of this node with parameterss renamed. If implemented on a parameter, this should make a copy of the parameter with a new name

Wraps a module to automatically apply `named_axes.nmap` to callables. The returned object will have similar attributes to `module`, including all

Constructs a named array of zeros with a given shape. Args: named_shape: Named shape for the result. dtype: Optional dtype for the result.