Autovisualizer

class penzai.treescope.autovisualize.Autovisualizer

Bases: Protocol

Protocol for autovisualizers.

An autovisualizer is a callable object that can be used to customize Treescope for rendering purposes. The active autovisualizer is called once for each subtree of a rendered PyTree, and can choose to:

• render it normally (but possibly customize its children),

• replace the rendering of that PyTree with a custom rich display object,

• or replace the current autovisualizer with a new one that will be active while rendering the children.

Methods

__init__(*args, **kwargs)
__call__(value, path)	Runs the autovisualizer on a node at a given path.

abstract __call__(value: Any, path: tuple[Any, ...] | None) → IPythonVisualization | CustomTreescopeVisualization | ChildAutovisualizer | None

Runs the autovisualizer on a node at a given path.

Parameters:

• value – A value being rendered in treescope.

• path – Path to this value from the root, as a JAX keypath. May be None if this object isn’t part of the root PyTree and so treescope doesn’t know how to access it.

Returns:

A visualization for this subtree, a child autovisualizer to use while rendering the node’s children, or None to continue rendering the node normally using the current autovisualizer.