class penzai.treescope.autovisualize.Autovisualizer[source]#

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.


__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[source]#

Runs the autovisualizer on a node at a given path.

  • 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.


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.