Treescope: an interactive HTML pretty-printer for ML notebooks

What problem Treescope solves, and for whom

Machine learning code produces deeply nested objects. A model is a dictionary of layers, each layer holds parameter arrays, and each array has a shape, a dtype and a distribution of values. The default IPython repr prints that as text, and once the nesting passes two or three levels the output becomes a wall of brackets that is hard to scan and impossible to navigate.

Treescope targets that specific moment. It is an interactive HTML pretty-printer and N-dimensional array visualizer built for machine learning and neural network research in IPython notebooks. The README describes it as a drop-in replacement for the standard IPython/Colab renderer. Its audience is narrow on purpose: people who inspect model internals from a notebook cell, not people who log metrics to a dashboard.

The feature list is practical rather than decorative. Subtrees expand and collapse so you can focus on one branch of a model. Arrays get faceted visualizations embedded directly in the output, so shape and value distribution are visible without a separate plotting call. Shared structures in a network are color-coded. A copy-path button on any node gives you the path to that part of the object, which is useful when you want to address it in later code. And the rendering strategy is customizable, so you can teach it about your own data structures.

The project started as the pretty-printer for the Penzai neural network library, and the README states it also supports networks built with Equinox, Flax NNX and PyTorch, plus basic JAX and NumPy code. That matters: you do not have to adopt Penzai to use Treescope.

How the rendering pipeline actually works

The mechanism is a rendering layer that sits between an object and the notebook's output area. Instead of letting IPython call repr() and print a string, Treescope walks the object graph, decides how to represent each node, and emits HTML with JavaScript attached for the interactive parts.

That walk is where the design choices show. Collapsing and expanding are per-node, so the renderer must keep a tree structure rather than a flat string. Arrays are not printed as numbers; they are passed to a visualizer that produces a faceted view, which is why the library declares numpy as its only required dependency and treats JAX and PyTorch as optional. Copy-path buttons imply the renderer tracks the access path to every node it emits, which is also what makes the path meaningful in later code.

There are two ways to put this in front of your output. The explicit route is treescope.show, which the README describes as being like print but producing a rich interactive output. The implicit route is treescope.register_as_default(), which installs Treescope as the notebook's default pretty printer so every cell output goes through it. The second is more convenient and more invasive: it changes rendering for everything, including objects you did not care about.

Array visualization is a separate switch. Calling treescope.active_autovisualizer.set_globally(treescope.ArrayAutovisualizer()) turns it on, and the combined helper treescope.basic_interactive_setup(autovisualize_arrays=True) does both jobs at once. The fact that autovisualization is opt-in rather than automatic is a deliberate trade-off: rendering every array in a large model as a faceted plot has a cost, and the library makes you ask for it.

One detail worth knowing is roundtrip mode. After rendering an object, clicking it and pressing the r key adds qualified names to every type in the visualization, which helps when the same structure appears under several types.

Install and first render in a notebook

Treescope is distributed on PyPI. The README gives a single install command, and the package requires Python 3.10 or newer according to pyproject.toml.

pip install treescope

After that, import it and either render one object explicitly or take over the notebook's default renderer. The README shows both entry points.

import treescope

treescope.show(my_model)

The output is an interactive HTML tree rather than text. Clicking the arrow buttons expands and collapses subtrees, and holding shift while scrolling scrolls horizontally instead of vertically.

If you want every cell to render through Treescope, register it as the default. This is the line that changes the notebook's behavior globally, so run it once near the top.

import treescope

treescope.register_as_default()

Array visualization is a second, separate step. The combined setup call enables both the default renderer and automatic array visualization.

import treescope

treescope.basic_interactive_setup(autovisualize_arrays=True)

If the output turns out too verbose, the README documents an abbreviation threshold that shortens collapsed objects at a given depth. It can be passed to the setup helper or set on treescope.abbreviation_threshold with set_globally or set_scoped.

import treescope

treescope.basic_interactive_setup(
    autovisualize_arrays=True,
    abbreviation_threshold=1,
)

For anything beyond these calls, the project points to its documentation site rather than the README.

Where Treescope stops being the right tool

The first limitation is structural: Treescope produces HTML. The interactivity described in the README, expanding subtrees, clicking copy-path buttons, pressing r for roundtrip mode, requires a browser rendering that HTML. In a terminal, a log file or a CI job, the interactive layer is inert. If your debugging happens in scripts and you read output in a shell, the default repr is not the problem you have.

The second is the cost of the default-renderer hook. register_as_default() changes how every object in the notebook is rendered, not just the model you are inspecting. That is convenient until it interacts with another library that also registers a pretty-printer, at which point you are debugging your debugging tool. The explicit treescope.show call avoids this entirely and is the safer starting point.

The third is array autovisualization at scale. Turning it on globally means every array that flows through the renderer gets a faceted visualization. For a small model that is helpful. For a large parameter tree it means a lot of rendering work for arrays you were not looking at. The README treats autovisualization as an explicit opt-in, and the abbreviation threshold exists precisely because verbose output is a real failure mode.

The fourth is scope. Treescope is a viewer. It does not train, profile, or track experiments, and the README makes no claim that it does. If you need to compare runs over time rather than inspect one object, this is the wrong category of tool.

Finally, the README carries the line that this is not an officially supported Google product. That is a statement about support expectations, not about code quality, but it belongs in the adoption decision.

Treescope against a plain repr or a debugger

The honest alternative for most people is the thing Treescope replaces: IPython's default repr, possibly with a custom __repr__ on your own classes. The difference is not cosmetic. A repr returns a string, which means nesting is flattened into text and there is no notion of a node you can collapse, no path you can copy, and no array visualization unless you wrote one. Treescope keeps the object graph intact through rendering, which is what makes the interactive features possible at all.

A second alternative is a step debugger such as pdb or the IPython debugger. That approach answers a different question. A debugger lets you pause execution and walk the live object graph at a breakpoint, which is more powerful for control-flow bugs. Treescope does not pause anything; it renders an object you already have, after the fact, in a cell. For inspecting a model's structure, the renderer is faster to reach for. For finding why a value became wrong three frames up, the debugger wins.

A third alternative is plotting arrays yourself with matplotlib. That gives you full control over the figure, and it is the right call when you need a specific chart for a paper. Treescope's array visualization is aimed at the quick look: shape and value distribution embedded in the same view as the surrounding structure, so you do not context-switch to a separate figure. The trade-off is that you get the visualization Treescope chooses, not one you designed.

None of these alternatives is strictly better. The choice depends on whether the object graph is the thing you are trying to read.

Maintenance, upgrades and licensing

The repository is not archived, and the last push was on 2026-06-18. Releases are infrequent: v0.1.8 in January 2025, v0.1.9 in February 2025, and v0.1.10 in August 2025. That cadence suggests a library that is feature-complete for its stated purpose rather than one under constant churn, but it also means fixes arrive on a slow schedule. Plan accordingly if you depend on a specific rendering behavior.

Upgrade cost looks low on the surface. The only required runtime dependency is numpy>=1.25.2, so installing Treescope does not drag a deep dependency tree into your environment. Python 3.10 or newer is required. The optional extras are where the weight sits: the test extra pulls in absl-py, jax, pytest, torch, pydantic and omegaconf, the notebook extra adds ipython, palettable and jax, and there are separate dev and docs extras. You only pay for those if you install them.

The version number is still 0.1.x. Semver does not bind pre-1.0 packages, so a minor bump can change behavior. Pin the version in your environment if rendered output matters to a shared notebook or a published figure.

On licensing: the package is Apache-2.0, declared both in the repository metadata and as a PyPI classifier. Apache-2.0 is a permissive license with an explicit patent grant and requires attribution and retention of notices. That is a summary of the license text, not legal advice; if you are redistributing Treescope inside a product, read the LICENSE file in the repository.

Editorial conclusion

Adopt Treescope if your debugging happens inside IPython or Colab and you regularly stare at nested parameter dictionaries, PyTree structures or model objects that the default repr flattens into unreadable text. Skip it if your workflow is script-first: the output is HTML, the interactivity lives in a browser, and a plain print or a logging call in a terminal gains nothing from it. Before committing, verify two things in your own environment: that treescope.register_as_default() does not clash with any other pretty-printer hook you already install, and that the array autovisualizer renders the dtypes and shapes you actually work with, not just float arrays.

Frequently asked questions

Official sources

1. google-deepmind/treescope on GitHub
2. License: Apache-2.0
3. Project website
4. README
5. Releases