Skip to main content

Treescope: An interactive HTML pretty-printer for ML research in IPython notebooks.

Project description

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:

Screenshot of a parameter dictionary without Treescope

You could be looking at this:

Screenshot of a parameter dictionary using Treescope

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

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

Getting Started

You can install Treescope using:

pip install treescope

and import it using:

import treescope

To render a specific object in an IPython notebook with Treescope, you can use treescope.show, which is like print but produces a rich interactive output.

Alternatively, you can simply configure Treescope as the default pretty printer for 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 and enable 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 the r key! This turns on "roundtrip mode", and adds qualified names to every type in the visualization, making it easier to identify what the types in your object are.

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

Looking for a neural network library with first-class support for Treescope's visualization features? Try Penzai!

Citation

If you have found Treescope to be useful for your research, please consider citing the following writeup (also available on arXiv):

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

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

treescope-0.1.6.tar.gz (530.1 kB view details)

Uploaded Source

Built Distribution

treescope-0.1.6-py3-none-any.whl (174.8 kB view details)

Uploaded Python 3

File details

Details for the file treescope-0.1.6.tar.gz.

File metadata

  • Download URL: treescope-0.1.6.tar.gz
  • Upload date:
  • Size: 530.1 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/5.1.1 CPython/3.12.5

File hashes

Hashes for treescope-0.1.6.tar.gz
Algorithm Hash digest
SHA256 67763ee24c841ada531521538a7dec407cb8902310634bead2e6bb315a5f2a2c
MD5 66baf08a259be19e5cc9393a9bc3ed24
BLAKE2b-256 e5787214ec914075d0f28b70719d3b684bfeb74e05a10394963feec9b4e86b78

See more details on using hashes here.

File details

Details for the file treescope-0.1.6-py3-none-any.whl.

File metadata

  • Download URL: treescope-0.1.6-py3-none-any.whl
  • Upload date:
  • Size: 174.8 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/5.1.1 CPython/3.12.5

File hashes

Hashes for treescope-0.1.6-py3-none-any.whl
Algorithm Hash digest
SHA256 a54cbe3cc5b5d9fcc0ac2c79dd3f2efe5f9c965d17e261e6d260a0fe13527a8f
MD5 ae6f9c788e277e8b239621e014e72f7e
BLAKE2b-256 1b4e3528887a1f79ccb8116f519a2f8cc289dcf75a215eb23da72930b280f4fe

See more details on using hashes here.

Supported by

AWS AWS Cloud computing and Security Sponsor Datadog Datadog Monitoring Fastly Fastly CDN Google Google Download Analytics Microsoft Microsoft PSF Sponsor Pingdom Pingdom Monitoring Sentry Sentry Error logging StatusPage StatusPage Status page