Processes nodes in the remainder to rebuild the tree.

Registers treescope as the default IPython renderer. This is a compatibility wrapper for `treescope.register_as_default`, which also turns on sup

Decorator to register a treescope-rendering color for a given effect. Args: color: The color to assign, as a CSS string. Returns: A deco

Renders a short summary of a struct. Can be used by other handlers that manipulate structs. Args: the_struct: Struct to render. Returns:

Interleaves intermediate values into a model. Args: root: The layer or model to annotate. argument: An argument to run the model with.

Returns a copy of `model` with `SaveIntermediate` layers inserted. Args: model: The model to copy. intermediates_group: The name of the gro

Scan a function over a named array axis while carrying along state. This function wraps `jax.lax.scan` to allow scanning over a named axis instead

Context manager for using scoped auto-generated StateVariable labels. Within this context manager, any StateVariable that does not have an explicit

Initializes a model, with constants and variables sharded based on a mesh. Args: initializer: The initializer to call. Should return a PyTree o

Shows a list of objects inline, like python print, but with rich display. This is a compatibility wrapper for `treescope.show`, which also formats

Renders the selection in IPython. This method is intended to visualize the selection object itself, and renders boxes around the selected nod

Renders an array of tokens. Args: tokens: An array of token IDs. vocab: Vocabulary, usually a SentencePiece tokenizer. axis_name_order:

Renders an array of token scores, but with the scored tokens as words. Args: tokens: An array of token IDs. scores: An array of token score

Renders the original tree in IPython, expanding up to the selected nodes. This method is intended to visualize a value but emphasizing the select

Stacks a sequence of named arrays along a named axis. Args: arrays: Sequence of named arrays to stack, which must have the same shape. axis

Rewrites docstrings for objects defined in `penzai.pz` to add alias links.

Runs temperature sampling in a Python for loop. Args: model: The converted model we are running inference with. prompt: A named array of pr

Runs temperature sampling in a Python for loop. Args: model: The converted model we are running inference with. initial_cache_state: The in