Repository files navigation

Treescope is an interactive HTML pretty-printer and N-dimensional array("tensor") visualizer, designed for machine learning and neural networksresearch in IPython notebooks. It's a drop-in replacement for the standardIPython/Colab renderer, and adds support for:

• Expanding and collapsing subtrees of rendered objects, to let you focus onthe parts of your model that you care about,
• Automatically embedding faceted visualizations of arbitrary-dimensional arraysand tensors directly into the output renderings, so you can quickly understandtheir 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 ofa rendered object,
• Customizing the visualization strategy to support rendering your own datastructures,
• And more!

Treescope was originally developed as the pretty-printer for thePenzai neural network library, butit also supports rendering neural networks developed with other libraries,includingEquinox,Flax NNX,andPyTorch.You can also use it with basicJAXandNumpycode.

With Treescope, instead of looking at this:

You could be looking at this:

This is an interactive visualization; try clicking the▶ buttons to expandparts of the output! (You can also hold shift while scrolling toscroll horizontally instead of vertically.)

Documentation on Treescope can be found athttps://treescope.readthedocs.io.

You can install Treescope using:

pip install treescope

and import it using:

importtreescope

To render a specific object in an IPython notebook with Treescope, you can usetreescope.show, which is likeprint but produces a rich interactive output.

Alternatively, you can simply configure Treescope as the default pretty printerfor your notebook via:

treescope.register_as_default()

To turn on automatic array visualization, you can run:

treescope.active_autovisualizer.set_globally(treescope.ArrayAutovisualizer())

Or, if you'd like to both set up Treescope as the default pretty printer andenable automatic array visualization, you can simply run:

treescope.basic_interactive_setup(autovisualize_arrays=True)

Once you've rendered an object, try clicking on it and pressing ther key!This turns on "roundtrip mode", and adds qualified names to every type in thevisualization, making it easier to identify what the types in your object are.

treescope.basic_interactive_setup(autovisualize_arrays=True,abbreviation_threshold=1,# or a different value)

For more information on how to use Treescope, check out theTreescope documentation.

Looking for a neural network library with first-class support for Treescope'svisualization features?TryPenzai!

If you have found Treescope to be useful for your research, please considerciting the following writeup (also available onarXiv):

@article{johnson2024penzai,    author={Daniel D. Johnson},    title={{Penzai} + {Treescope}: A Toolkit for Interpreting, Visualizing, and Editing Models As Data},    year={2024},    journal={ICML 2024 Workshop on Mechanistic Interpretability}}

This is not an officially supported Google product.