treescope

Treescope is an interactive HTML pretty-printer and ND-array/tensor visualiser. Originally designed for penzai neural networks in jax.

• show shapes and the distribution of their values NDArray
• colour-coding parts of neural network models to emphasise shared structures

render:

• numpy/jax/pytorch NDArray/tensor or penzai NamedArray
• neural network models
  • penzai / equinox pytree dataclasses
  • nnx nnx.display(model)
  • pytorch dynamic Python objects
• dicts, lists, tuples, and sets
• dataclasses and namedtuples
• functions
• builtins and literals
• arbitrary pytree types

visualise some arrays

with treescope.active_autovisualizer.set_scoped(
	treescope.ArrayAutovisualizer()
):
	treescope.display(arrays) # inspect nested object containing NDArrays
	treescope.render_array(arr) # NDArray visualizer

set as default ipython pretty-printer

treescope.basic_interactive_setup()

render to html

with treescope.active_autovisualizer.set_scoped(
	treescope.ArrayAutovisualizer()
):
	contents = treescope.render_to_html(some_arrays)
 
with open("/tmp/treescope_output.html", "w") as f:
	f.write(contents)