Treescope

0.1.10 · active · verified Fri Apr 10

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 serves as a drop-in replacement for the standard IPython/Colab renderer, enhancing output with features like expandable/collapsible subtrees, embedded faceted visualizations of arbitrary-dimensional arrays, and color-coding for model structures. It supports inspecting data from libraries such as JAX, NumPy, PyTorch, Equinox, Flax NNX, and Penzai. Maintained by Google DeepMind, the current version is 0.1.10.

Warnings

gotcha Treescope's rich HTML output and automatic rendering features are only active in IPython/Colab notebooks after `treescope.basic_interactive_setup()` or manual registration. Without this setup, objects will revert to standard `repr()` output.
Fix: Call `treescope.basic_interactive_setup()` at the beginning of your notebook session to enable Treescope as the default renderer and activate array autovisualization.
gotcha The full interactive power of Treescope (e.g., expanding/collapsing subtrees, 'roundtrip mode' with the 'r' key for qualified names, 'copy path' buttons) requires direct user interaction within the rendered HTML output in an IPython environment, which might not be immediately obvious to new users.
Fix: Encourage users to click on rendered objects and experiment with keyboard shortcuts (like 'r') within the IPython output to unlock interactive exploration. Refer to official documentation for interactive features.
gotcha While Treescope supports common Python types and popular ML libraries, visualizing highly custom or complex data structures may require implementing custom `treescope` rendering logic to achieve optimal interactive display.
Fix: For bespoke data structures, consult the 'Customizing Treescope' section of the official documentation to learn how to define custom visualizers or `__treescope_repr__` methods.

Install

pip install treescope Install stable version

Imports

treescope
```
import treescope
```
show
```
treescope.show(...)
```
Used to explicitly render objects with Treescope, similar to `print`.
display
```
treescope.display(...)
```
Used to display a value as an interactively foldable object.
basic_interactive_setup
```
treescope.basic_interactive_setup()
```
Calling `basic_interactive_setup()` is crucial for Treescope to act as the default IPython renderer and enable automatic array visualization.

Quickstart

This quickstart demonstrates how to enable Treescope as the default renderer in an IPython/Colab environment and visualize a complex, nested data structure containing both JAX and NumPy arrays. After `basic_interactive_setup()`, simply outputting a variable in a notebook cell will automatically trigger Treescope's rich, interactive HTML visualization.

import treescope
import jax
import jax.numpy as jnp
import numpy as np

# Enable Treescope as the default IPython renderer and auto-visualizers
treescope.basic_interactive_setup()

# Define a nested data structure with arrays
my_data = {
    "model_config": {
        "num_layers": 4,
        "hidden_size": 256,
        "activation": "relu"
    },
    "weights": {
        "layer_0": jnp.ones((16, 256)),
        "layer_1": jnp.zeros((256, 128))
    },
    "metadata": [
        "training_started",
        "epoch_1",
        {"loss": 0.123, "accuracy": 0.98, "grad_norm": np.array(0.005)}
    ]
}

# In an IPython notebook, simply outputting the variable will use Treescope:
my_data

# Or explicitly display it:
treescope.display(my_data)

view raw JSON →