Apply RoPE to a one-dimensional JAX array.

Apply RoPE to a one-dimensional JAX array.

Constructs a zeros array with a reasonable dtype and a given structure.

Zeros initializer for named arrays.

Zeros initializer for named arrays.

Default predicate that selects only untagged holes.

Pretty-prints this object for an IPython pretty-printer.

Implementation of the training step for StatefulTrainer.

Custom wrapper hook that intercepts selected nodes.

Prepends a prefix to all parameter names inside the tree. This function can be used to avoid name conflicts when combining sublayers into a large

Annotates shapes in a model or layer, for inputs with this input structure. Args: root: The layer or model to annotate. dummy_input: A dumm

Convenience function to create a range along a named axis. This is shorthand for ``wrap(jnp.arange(...)).tag(name)``. Args: name: Name for t

Abstracts a concrete tree of values into a tree of `ArraySpec` nodes. Args: tree: The tree to get an array structure for. Returns: A cop

Initializes the parameter, returning an equivalent frozen one.

Asserts that the given model has no ParameterSlot subtrees.

Helper property for index update functionality. Lifts the ``jax.Array.at[...]`` syntax to also work for Penzai NamedArrays. Similar to direc

Wraps a submodel in a handler that takes ownership of shareable params. Penzai models must have only one copy of each parameter in the model pytree

Sets up IPython for interactive use with Treescope. This is a compatibility wrapper for `treescope.basic_interactive_setup`, which also turns on

Collects effect types of each broken `HandledEffectRef` in a (sub)model. A broken handler reference occurs when a model contains a `HandledEffectRe

Builds a train step function for a common training loop. For simplicity, the output of the train step function is the third output of the loss fu

Calls a model with an argument, and captures the selected submodel. This function is designed to enable quickly exploring the behavior of small p

Calls a model with an argument, and captures the selected submodel. This function is designed to enable quickly exploring the behavior of small p

Captures positional and keyword arguments from this call.

Checks that a layer has been configured correctly by tracing it. This function runs the layer under `jax.eval_shape` and passes in a dummy input

Checks that a structure of values matches a pattern. Args: value: A PyTree of arrays to check. pattern: A PyTree of `ArraySpec` leaves to c

Decorator for ``Layer.__call__`` to add shape-checking and name scopes.

Combines leaves from multiple partitions. This function can be used to reverse the results of calling `partition` on a selector. It walks two or

Computes loss, aux out, and loss state for a single example.

Computes loss, aux out, and state updates for a single example.

Runs a loss function and computes all of its outputs. This function runs the model and loss function and updates the corresponding parameters in

Concatenates a sequence of named arrays along a named axis. Args: arrays: Sequence of named arrays to concatenate, which must have the same

Returns a disabled version of this layer.

Returns a disabled version of this layer.

Raises an error, since disabling is no longer supported. Originally, global overrides for Penzai/Treescope ContextualValue objects needed to be t

Displays a value as an interactively foldable object. This is a compatibility wrapper for `treescope.display`, which also formats objects impleme

Displays a top-level selection object in IPython. This function produces a rendering of either the selection itself, or its selected value. Howev

Custom displayhook to display results with IPython.

The dtype of the wrapped array.