penzai.treescope: Penzai’s interactive pretty-printer

`penzai.treescope`: Penzai’s interactive pretty-printer#

You can configure treescope as the default IPython pretty-printer using

pz.ts.basic_interactive_setup()

or, for more control:

pz.ts.register_as_default()
pz.ts.register_autovisualize_magic()
pz.enable_interactive_context()
pz.ts.active_autovisualizer.set_interactive(pz.ts.ArrayAutovisualizer())

You can also pretty-print individual values using pz.show or pz.ts.display. See also the pz.ts module for convenient access to the most commonly-used treescope functions and classes.

Extending and Customizing Treescope#

Adding Treescope Support for New Types#

To teach Treescope how to render your custom type, you can define the special method __penzai_repr__ with the following signature:

NodePath = <implementation detail>
Rendering = <implementation detail>

class MyCustomType:
  ...

  def __penzai_repr__(
    self,
    path: NodePath,
    subtree_renderer: Callable[[Any, NodePath], Rendering],
  ) -> Rendering | type(NotImplemented):
    ...

The types Treescope uses to represent node paths and renderings should be considered internal implementation details and may change in the future. To ensure your implementation of __penzai_repr__ is compatible with future changes, you can use the following high-level API to build renderings:

`repr_lib.render_object_constructor`(...[, ...])	Renders an object in "constructor format", similar to a dataclass.
`repr_lib.render_dictionary_wrapper`(...[, ...])	Renders an object in "wrapped dictionary format".

This API will be expanded in the future. If you would like to customize the rendering of your type beyond what is currently possible with repr_lib, please open an issue!

Customizing Automatic Visualization#

By default, Treescope will automatically visualize NDArrays using its own array visualizer. However, you can hook into this system to insert arbitrary figure objects into the leaves of your tree, replacing arbitrary objects.

To do this, you can define an autovisualization function:

def autovisualizer_fn(
    value: Any,
    path: NodePath | None,
) -> pz.ts.IPythonVisualization | pz.ts.ChildAutovisualizer | None:
  ...

For instance, to render certain objects with Plotly, you could define

def my_plotly_autovisualizer(
    value: Any,
    path: NodePath | None,
):
  if isinstance(value, (np.ndarray, jax.Array)):
    return pz.ts.IPythonVisualization(
        px.histogram(
            value.flatten(),
            width=400, height=200
        ).update_layout(
            margin=dict(l=20, r=20, t=20, b=20)
        )
    )

You can then enable your autovisualizer in a scope:

with pz.ts.active_autovisualizer.set_scoped(
    my_plotly_autovisualizer
):
    IPython.display.display(some_object_with_arrays)

Alternatively, you can also pass custom visualizers to the %%autovisualize magic at the top of an IPython cell to let it handle the set_scoped boilerplate for you:

%%autovisualize my_plotly_autovisualizer
some_object_with_arrays

You can also disable automatic visualization entirely using %%autovisualize None.

Advanced: Custom Handlers and Intermediate Representation (Unstable)#

Advanced users who wish to fully customize how Treescope builds renderings can construct their own renderer objects and handlers, or directly construct renderings using Treescope’s intermediate representation. Renderer objects and the expected types of handlers are defined in penzai.treescope.renderer, and the intermediate representation is currently defined in penzai.treescope.foldable_representation.

Warning

The Treescope intermediate representation and handler system will be changing in the future. We recommend only using this API for experimentation, and not using it in library code. Code that uses the internals of Penzai’s intermediate representation may break at any time, and should pin a specific version of Penzai to avoid future changes.

Other Treescope Methods#

The following methods should usually be accessed using the pz.ts alias module instead of accessing them directly from treescope.

Using Treescope in IPython Notebooks#

`treescope_ipython.register_as_default`([...])	Registers treescope as the default IPython renderer.
`treescope_ipython.register_autovisualize_magic`()	Registers the `%%autovisualize` magic.
`treescope_ipython.register_context_manager_magic`()	Registers the `%%with` magic.

Showing Objects Explicitly#

`treescope_ipython.display`(value[, ...])	Displays a value as an interactively foldable object.
`treescope_ipython.show`(*args[, wrap, ...])	Shows a list of objects inline, like python print, but with rich display.
`arrayviz.arrayviz.render_array`(array, *[, ...])	Renders an array (positional or named) to a displayable HTML object.
`arrayviz.arrayviz.render_array_sharding`(array)	Renders the sharding of an array.
`arrayviz.arrayviz.render_sharded_shape`(...)	Renders the sharding an array would have, based on its shape.
`arrayviz.arrayviz.integer_digitbox`(value[, ...])	Returns a "digitbox" rendering of a single integer.
`arrayviz.arrayviz.text_on_color`(text, value)	Renders some text on colored background, similar to arrayviz coloring.

Styling Displayed Objects#

`figures.inline`(*parts[, wrap])	Returns a figure that arranges a set of displayable objects along a line.
`figures.indented`(subfigure)	Returns a figure object that displays a value with an indent.
`figures.with_font_size`(subfigure, size)	Returns a scaled version of the first figure.
`figures.with_color`(subfigure, color)	Returns a colored version of the first figure.
`figures.bolded`(subfigure)	Returns a bolded version of the first figure.
`figures.styled`(subfigure, style)	Returns a CSS-styled version of the first figure.

Configuring Treescope#

`default_renderer.active_renderer`	The default renderer to use when rendering a tree to HTML.
`default_renderer.active_expansion_strategy`	The default expansion strategy to use when rendering a tree to HTML.
`default_renderer.using_expansion_strategy`([...])	Context manager that modifies the expansion strategy of Treescope.
`autovisualize.active_autovisualizer`	The active autovisualizer to use when rendering a tree to HTML.
`arrayviz.arrayviz.default_diverging_colormap`	Default diverging colormap.
`arrayviz.arrayviz.default_sequential_colormap`	Default sequential colormap.

Building Autovisualizers#

`autovisualize.Autovisualizer`	Protocol for autovisualizers.
`autovisualize.ChildAutovisualizer`	Used by autovisualizers to switch to a different autovisualizer.
`autovisualize.IPythonVisualization`	Used by autovisualizers to replace a subtree with a display object.
`autovisualize.CustomTreescopeVisualization`	Used by autovisualizers to directly insert Treescope content.
`treescope_ipython.default_magic_autovisualizer`	The default autovisualizer to use for the `%%autovisualize` magic.
`arrayviz.array_autovisualizer.ArrayAutovisualizer`	An automatic visualizer for arrays.

Rendering to Strings#

`default_renderer.render_to_text`(value[, ...])	Renders an object to text using the default renderer.
`default_renderer.render_to_html`(value[, ...])	Renders an object to HTML using the default renderer.

Utility Types#

copypaste_fallback.NotRoundtrippable

A placeholder for a non-roundtrippable object in roundtrip mode.