Treescope

Treescope is an interactive HTML pretty-printer and N-dimensional array (“tensor”) visualizer, designed for machine learning and neural networks research in IPython notebooks. It’s a drop-in replacement for the standard IPython/Colab renderer, and adds support for:

• Expanding and collapsing subtrees of rendered objects, to let you focus on the parts of your model that you care about,

• Automatically embedding faceted visualizations of arbitrary-dimensional arrays and tensors directly into the output renderings, so you can quickly understand their shapes and the distribution of their values,

• Color-coding parts of neural network models to emphasize shared structures,

• Inserting “copy path” buttons that let you easily copy the path to any part of a rendered object,

• Customizing the visualization strategy to support rendering your own data structures,

• And more!

Treescope was originally developed as the pretty-printer for the Penzai neural network library, but it also supports rendering neural networks developed with other libraries, including Equinox, Flax NNX, and PyTorch. You can also use it with basic JAX and Numpy code.

With Treescope, instead of looking at this: