### Install Treescope

Source: https://github.com/google-deepmind/treescope/blob/main/README.md

Install the Treescope library using pip.

```bash
pip install treescope
```

--------------------------------

### Basic Interactive Setup for Treescope

Source: https://github.com/google-deepmind/treescope/blob/main/README.md

Perform a basic setup for Treescope, including registering it as the default pretty printer and enabling automatic array visualization.

```python
treescope.basic_interactive_setup(autovisualize_arrays=True)
```

--------------------------------

### Install and Import Treescope

Source: https://github.com/google-deepmind/treescope/blob/main/docs/index.rst

Install the package via pip and import it into your environment.

```bash
pip install treescope
```

```python
import treescope
```

--------------------------------

### Setup Basic Interactive Mode

Source: https://github.com/google-deepmind/treescope/blob/main/docs/notebooks/building_custom_visualizations.ipynb

Enables basic interactive setup for Treescope, which is necessary for its interactive features to work.

```python
import treescope
treescope.basic_interactive_setup()
```

--------------------------------

### Basic array visualization examples

Source: https://github.com/google-deepmind/treescope/blob/main/docs/notebooks/building_custom_visualizations.ipynb

Examples of default array visualization behavior in Treescope.

```python
np.arange(10)
```

```python
[np.arange(10), np.linspace(-10,10,20)]
```

--------------------------------

### Install Treescope in Notebooks

Source: https://github.com/google-deepmind/treescope/blob/main/docs/notebooks/array_visualization.ipynb

Checks for the treescope package and installs it if missing, suitable for Colab or Kaggle environments.

```python
try:
  import treescope
except ImportError:
  !pip install treescope
```

--------------------------------

### Basic Interactive Setup with Autovisualization

Source: https://context7.com/google-deepmind/treescope/llms.txt

Sets up Treescope as the default renderer in IPython, enabling automatic visualization of arrays. Configure abbreviation thresholds to control verbosity.

```python
import treescope
import numpy as np

# Full interactive setup with array autovisualization enabled
treescope.basic_interactive_setup(autovisualize_arrays=True)

# Optional: Set abbreviation threshold to reduce verbosity
treescope.basic_interactive_setup(
    autovisualize_arrays=True,
    abbreviation_threshold=1,  # Abbreviate collapsed objects at depth 1
)

# Now any object displayed in the notebook will use Treescope
data = {
    "weights": np.random.randn(10, 5),
    "biases": np.zeros(5),
    "config": {"learning_rate": 0.001, "epochs": 100}
}
# Simply evaluate or use display() - Treescope renders automatically
data
```

--------------------------------

### Configuration and Setup

Source: https://github.com/google-deepmind/treescope/blob/main/docs/api/treescope.rst

Functions to configure Treescope as the default renderer in IPython and manage global settings.

```APIDOC
## Configuration and Setup

### Description
Methods to enable Treescope as the default renderer and register IPython magics.

### Functions
- **basic_interactive_setup()** - Configures Treescope as the default renderer for all IPython cell outputs.
- **register_as_default()** - Registers Treescope as the default renderer.
- **register_autovisualize_magic()** - Registers the %%autovisualize magic.
- **register_context_manager_magic()** - Registers the %%with magic.
```

--------------------------------

### Instantiate Custom Type

Source: https://github.com/google-deepmind/treescope/blob/main/docs/notebooks/building_custom_visualizations.ipynb

Example of how to instantiate a custom type that has implemented `__treescope_repr__`.

```python
MySimpleType(123)
```

--------------------------------

### Compose and Style Figures with Treescope

Source: https://context7.com/google-deepmind/treescope/llms.txt

Demonstrates using the figures module to compose multiple display objects into figures. Includes examples for inline display, word wrapping, indentation, CSS styling, font adjustments, coloring, bolding, and value-based text coloring.

```python
import treescope
from treescope import figures
import numpy as np

treescope.register_as_default()

# Inline multiple objects horizontally
fig = figures.inline("Label:", np.array([1, 2, 3]), "more text")
treescope.display(fig)

# With word wrapping
fig = figures.inline("A", "B", "C", "D", wrap=True)
treescope.display(fig)

# Indented content
fig = figures.indented({"nested": [1, 2, 3]})
treescope.display(fig)

# Apply CSS styling
fig = figures.styled(
    {"data": [1, 2, 3]},
    "background-color: #f0f0f0; padding: 10px;"
)
treescope.display(fig)

# Font size adjustments
fig = figures.with_font_size({"small": "text"}, size=0.8)  # 80% of normal
fig = figures.with_font_size({"large": "text"}, size="18pt")  # Explicit size
treescope.display(fig)

# Color text
fig = figures.with_color("Warning message", color="red")
treescope.display(fig)

# Bold text
fig = figures.bolded("Important!")
treescope.display(fig)

# Text on colored background based on value
fig = figures.text_on_color("high", value=0.9, vmax=1.0)
treescope.display(fig)

fig = figures.text_on_color("negative", value=-0.5, vmin=-1.0, vmax=1.0)
treescope.display(fig)
```

--------------------------------

### Enable Treescope Features

Source: https://context7.com/google-deepmind/treescope/llms.txt

Call this at the start of a notebook session to enable all Treescope features, including automatic visualization of arrays.

```python
import treescope
treescope.basic_interactive_setup(autovisualize_arrays=True)
```

--------------------------------

### Define Dataclasses and Custom Repr for Treescope

Source: https://github.com/google-deepmind/treescope/blob/main/docs/notebooks/pretty_printing.ipynb

Setup custom classes to demonstrate nested object visualization and multi-line representation.

```python
import dataclasses

@dataclasses.dataclass
class MyDataclass:
  a: Any
  b: Any
  c: Any

class TheZenOfPython:
  def __repr__(self):
    return "<The Zen of Python:\nBeautiful is better than ugly.\nExplicit is better than implicit.\nSimple is better than complex.\nComplex is better than complicated.\nFlat is better than nested.\nSparse is better than dense.\nReadability counts.\nSpecial cases aren't special enough to break the rules.\nAlthough practicality beats purity.\nErrors should never pass silently.\nUnless explicitly silenced.\nIn the face of ambiguity, refuse the temptation to guess.\nThere should be one-- and preferably only one --obvious way to do it.\nAlthough that way may not be obvious at first unless you're Dutch.\nNow is better than never.\nAlthough never is often better than *right* now.\nIf the implementation is hard to explain, it's a bad idea.\nIf the implementation is easy to explain, it may be a good idea.\nNamespaces are one honking great idea -- let's do more of those!>"
```

```python
[
    MyDataclass('a' * i, 'b' * i, ('cccc\n') * i)
    for i in range(10)
] + [
    MyDataclass(TheZenOfPython(), TheZenOfPython(), TheZenOfPython())
]
```

--------------------------------

### Configure Treescope Expansion Strategy

Source: https://context7.com/google-deepmind/treescope/llms.txt

Utilizes a context manager to configure how Treescope expands and collapses nodes in the output layout. Examples show limiting height, controlling width, and fine-grained expansion behavior.

```python
import treescope
import numpy as np

treescope.register_as_default()

nested_data = {
    "level1": {
        "level2": {
            "level3": {"values": list(range(100))}
        }
    }
}

# Limit maximum height of expanded output
with treescope.using_expansion_strategy(max_height=10):
    treescope.display(nested_data)

# Allow unlimited height, control by width only
with treescope.using_expansion_strategy(max_height=None, target_width=80):
    treescope.display(nested_data)

# Fine-grained control over expansion behavior
with treescope.using_expansion_strategy(
    max_height=20,
    target_width=60,
    relax_for_only_child=True,   # Allow exceeding height for single children
    recursive_expand_height=7,    # Pre-expand collapsed node children
):
    treescope.display(nested_data)
```

--------------------------------

### Render Data to Text with Roundtrip Mode

Source: https://context7.com/google-deepmind/treescope/llms.txt

Renders data to text, with an option to enable roundtrip mode for preserving formatting. Includes an example of ignoring exceptions during rendering.

```python
text = treescope.render_to_text(data, roundtrip_mode=True)
print(text)
```

```python
text = treescope.render_to_text(data, ignore_exceptions=True)
```

--------------------------------

### Render Array with Multiple Sliders

Source: https://github.com/google-deepmind/treescope/blob/main/docs/notebooks/array_visualization.ipynb

Visualize a complex array with sliders on multiple axes simultaneously. This example uses sliders for 'row_wavelength' and 'col_wavelength'.

```python
row_wavelength = (4 * jnp.arange(10) + 4)[:, None, None, None]
col_wavelength = (4 * jnp.arange(10) + 4)[None, :, None, None]
col = np.arange(15)[None, None, :, None]
row = np.arange(15)[None, None, None, :]

values = (
    jnp.sin(2 * np.pi * row / row_wavelength)
    * jnp.sin(2 * np.pi * col / col_wavelength)
)
```

```python
treescope.render_array(
    values,
    columns=[2],
    sliders=[0, 1],
    axis_item_labels={
        0: [str(v) for v in row_wavelength.squeeze((1, 2, 3))],
        1: [str(v) for v in col_wavelength.squeeze((0, 2, 3))],
    },
    axis_labels={
        0: "row_wavelength",
        1: "col_wavelength",
        2: "row",
        3: "col",
    }
)
```

--------------------------------

### treescope.basic_interactive_setup

Source: https://context7.com/google-deepmind/treescope/llms.txt

Configures IPython for interactive use with Treescope, setting it as the default renderer and enabling optional array visualization.

```APIDOC
## treescope.basic_interactive_setup

### Description
Sets up IPython for interactive use with Treescope, configuring it as the default renderer, registering cell magics, and optionally enabling automatic array visualization.

### Parameters
#### Query Parameters
- **autovisualize_arrays** (bool) - Optional - Whether to enable automatic array visualization.
- **abbreviation_threshold** (int) - Optional - Depth at which to abbreviate collapsed objects.
```

--------------------------------

### Create and Shard a JAX Array

Source: https://github.com/google-deepmind/treescope/blob/main/docs/notebooks/array_visualization.ipynb

Demonstrates how to create a device mesh and shard a JAX array across multiple devices using `PositionalSharding`.

```python
from jax.experimental import mesh_utils
devices = mesh_utils.create_device_mesh((8,))
pos_sharding = jax.sharding.PositionalSharding(devices).reshape((4, 2))

sharded_array = jax.device_put(
    jnp.arange(512).reshape((16, 32)), pos_sharding
)
```

--------------------------------

### Render Array with Single Slider

Source: https://github.com/google-deepmind/treescope/blob/main/docs/notebooks/array_visualization.ipynb

Use `render_array` to visualize an array with interactive sliders. This example shows a single slider for the 'time' axis.

```python
treescope.render_array(
    values_over_time,
    columns=[1],
    sliders=[0],
    axis_labels={
        0: "time",
        1: "row",
        2: "col",
    }
)
```

--------------------------------

### Import Dependencies

Source: https://github.com/google-deepmind/treescope/blob/main/docs/notebooks/array_visualization.ipynb

Standard imports for JAX, Numpy, and IPython environments.

```python
from __future__ import annotations
from typing import Any

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

import IPython
```

--------------------------------

### Using the %%autovisualize magic

Source: https://github.com/google-deepmind/treescope/blob/main/docs/notebooks/building_custom_visualizations.ipynb

Applying an autovisualizer using the IPython magic command to avoid manual scoping boilerplate.

```python
%%autovisualize my_plotly_autovisualizer
{
    "foo": np.arange(10)[:, None] * np.arange(10)[None, :],
    "bar": np.sin(np.arange(100) * 0.1).reshape((10,10))
}
```

--------------------------------

### Access help for render_array

Source: https://github.com/google-deepmind/treescope/blob/main/docs/notebooks/array_visualization.ipynb

Displays the docstring for the render_array function.

```python
help(treescope.render_array)
```

--------------------------------

### treescope.show

Source: https://context7.com/google-deepmind/treescope/llms.txt

Displays multiple objects inline with rich HTML output, similar to Python's print function.

```APIDOC
## treescope.show

### Description
Displays multiple objects inline with rich HTML output, similar to Python's print but with interactive visualization capabilities.

### Parameters
#### Query Parameters
- **space_separated** (bool) - Optional - Whether to include spaces between items.
- **wrap** (bool) - Optional - Enable word wrapping at line ends.
- **autovisualize** (bool) - Optional - Force array autovisualization for this call.
```

--------------------------------

### Display Objects with Treescope

Source: https://context7.com/google-deepmind/treescope/llms.txt

Use `treescope.show()` for quick inline displays or `treescope.display()` for detailed exploration of individual objects.

```python
treescope.show(my_object)
```

```python
treescope.display(my_object)
```

--------------------------------

### Visualize Dataclasses

Source: https://github.com/google-deepmind/treescope/blob/main/docs/notebooks/pretty_printing.ipynb

Renders dataclass instances, including those with custom initialization logic via roundtrip mode.

```python
@dataclasses.dataclass(frozen=True)
class Bar:
  c: str
  d: int
  some_list: list = dataclasses.field(default_factory=list)

IPython.display.display(Bar(c="bar", d=2))
```

```python
@dataclasses.dataclass
class WeirdInitClass:
  foo: int

  def __init__(self, half_foo: int):
    self.foo = 2 * half_foo

# This shows as WeirdInitClass(foo=4):
treescope.display(WeirdInitClass(2))

# But in roundtrip mode (explicit or after pressing `r`), it shows as
#   pz.dataclass_from_attributes(WeirdInitClass, foo=4)
# which bypasses __init__ and rebuilds the dataclass's attributes directly,
# since __init__ doesn't take `foo` as an argument.
treescope.display(WeirdInitClass(2), roundtrip_mode=True)
```

--------------------------------

### Configure Global Autovisualizer

Source: https://context7.com/google-deepmind/treescope/llms.txt

Sets a global autovisualizer for automatic array rendering and token lookup visualization.

```python
treescope.active_autovisualizer.set_globally(autovisualizer)

# Arrays are now automatically visualized with these settings
weights = np.random.randn(64, 32)
treescope.display(weights)

# Create an autovisualizer for tokenized text with token lookup
def token_lookup(token_id):
    vocab = {0: "<pad>", 1: "<unk>", 2: "hello", 3: "world"}
    return vocab.get(token_id, f"[{token_id}]")

tokenizer_viz = treescope.ArrayAutovisualizer.for_tokenizer(token_lookup)
treescope.active_autovisualizer.set_globally(tokenizer_viz)

# Token arrays now show token strings on hover
tokens = np.array([[2, 3, 0], [3, 2, 1]])
treescope.display(tokens)
```

--------------------------------

### Configure Array Autovisualizer

Source: https://context7.com/google-deepmind/treescope/llms.txt

Creates a customized ArrayAutovisualizer with specific settings for array rendering, including maximum size, cutoff per axis, edge item display, centering around zero, pixel size, and preferred axis orientation.

```python
import treescope
import numpy as np

# Create a customized array autovisualizer
autovisualizer = treescope.ArrayAutovisualizer(
    maximum_size=4000,         # Max elements to show before truncation
    cutoff_size_per_axis=128,  # Max elements per axis
    edge_items=5,              # Values to keep when truncating
    around_zero=True,          # Center continuous data around zero
    pixels_per_cell=7,         # Pixel size for each array element (1-21)
    prefers_column=("seq",),   # Axis names to prefer as columns
    prefers_row=("batch",),    # Axis names to prefer as rows
)

```

--------------------------------

### Initialize array for slider scrubbing

Source: https://github.com/google-deepmind/treescope/blob/main/docs/notebooks/array_visualization.ipynb

Prepares an array for interactive slicing using the sliders argument.

```python
time = jnp.arange(100)[:, None, None]
col = np.linspace(-2, 2, 15)[None, :, None]
row = np.linspace(-2, 2, 15)[None, None, :]

values_over_time = jax.nn.sigmoid(
    0.05 * time - 2 - row - jnp.sin(2 * col - 0.1 * time)
)
```

--------------------------------

### Import necessary libraries

Source: https://github.com/google-deepmind/treescope/blob/main/docs/_include/_glue_figures.ipynb

Imports required libraries for notebook integration, JAX, and Treescope.

```python
import os
import functools

import myst_nb

import treescope

import jax
import jax.numpy as jnp

import IPython.utils.capture

from penzai import pz
```

--------------------------------

### Implement Custom Type Rendering with __treescope_repr__

Source: https://context7.com/google-deepmind/treescope/llms.txt

Shows how to define a `__treescope_repr__` method on custom classes to enable Treescope to render them with specific formatting. Uses `repr_lib.render_object_constructor` for structured output.

```python
import treescope
from treescope import repr_lib

treescope.register_as_default()

class NeuralLayer:
    def __init__(self, in_features, out_features, activation="relu"):
        self.in_features = in_features
        self.out_features = out_features
        self.activation = activation

    def __treescope_repr__(self, path, subtree_renderer):
        return repr_lib.render_object_constructor(
            object_type=type(self),
            attributes={
                "in_features": self.in_features,
                "out_features": self.out_features,
                "activation": self.activation,
            },
            path=path,
            subtree_renderer=subtree_renderer,
            roundtrippable=True,  # Can reconstruct from repr
            color="#e6f3ff",      # Light blue background
        )

# Now displays with custom formatting
layer = NeuralLayer(128, 64, "gelu")
treescope.display(layer)
```

--------------------------------

### Render a basic array

Source: https://github.com/google-deepmind/treescope/blob/main/docs/notebooks/array_visualization.ipynb

Creates and renders a 2D numpy array using default visualization settings.

```python
my_array = np.cos(np.arange(300).reshape((10,30)) * 0.2)

treescope.render_array(my_array)
```

--------------------------------

### Visualize Neural Network Models

Source: https://github.com/google-deepmind/treescope/blob/main/docs/notebooks/pretty_printing.ipynb

Inspects and renders model architectures for Penzai, Equinox, and PyTorch.

```python
from penzai.models import simple_mlp
simple_mlp.MLP.from_config(
    name="mlp",
    init_base_rng=jax.random.key(0),
    feature_sizes=[8, 32, 32, 8]
)
```

```python
torch.nn.Sequential(
    torch.nn.Linear(8, 32),
    torch.nn.ReLU(),
    torch.nn.Linear(32, 32),
    torch.nn.ReLU(),
    torch.nn.Linear(32, 8),
)
```

--------------------------------

### Build and process a Llama-like Transformer model

Source: https://github.com/google-deepmind/treescope/blob/main/docs/_include/_glue_figures.ipynb

Builds a Llama-like transformer model using Penzai and processes its parameters for display.

```python
from penzai.models.transformer.variants import llamalike_common

model = llamalike_common.build_llamalike_transformer(
    llamalike_common.LlamalikeTransformerConfig(
        num_kv_heads=8,
        query_head_multiplier=1,
        embedding_dim=256,
        projection_dim=32,
        mlp_hidden_dim=512,
        num_decoder_blocks=10,
        vocab_size=1000,
        mlp_variant="geglu_approx",
        rope_wavelength=10_000,
        tie_embedder_and_logits=True,
        use_layer_stack=False,
        parameter_dtype=jnp.float32,
        activation_dtype=jnp.float32,
    ),
    init_base_rng=jax.random.key(42),
)

_, params = pz.unbind_params(model, freeze=True)
params = pz.select(params).at(lambda root: root[0].value).apply(
    lambda x: x.order_as("embedding", "vocabulary")
)
```

--------------------------------

### Configure XLA Device Count

Source: https://github.com/google-deepmind/treescope/blob/main/docs/notebooks/array_visualization.ipynb

Sets the XLA_FLAGS environment variable to force multiple host platform devices.

```python
import os
os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=8"
```

--------------------------------

### Treescope Object Documentation

Source: https://github.com/google-deepmind/treescope/blob/main/docs/_templates/pzbase.rst

Standard documentation format for Treescope modules and objects.

```APIDOC
## Treescope Object Reference

### Description
This section provides documentation for specific objects and modules within the Treescope library using the auto-documentation directive.

### Usage
Use the current module context followed by the auto-directive for the specific object type (e.g., class, function, attribute).

### Syntax
.. currentmodule:: [module_name]

.. auto[objtype]:: [objname]
```

--------------------------------

### Import Treescope

Source: https://github.com/google-deepmind/treescope/blob/main/README.md

Import the Treescope library for use in your Python code.

```python
import treescope
```

--------------------------------

### Render Array Sharding

Source: https://context7.com/google-deepmind/treescope/llms.txt

Visualizes how distributed arrays are partitioned across devices using JAX.

```python
import treescope
import numpy as np

# This example requires JAX with multiple devices
# import jax
# import jax.numpy as jnp
# from jax.sharding import NamedSharding, PartitionSpec, Mesh

# # Create a sharded array
# devices = jax.devices()
# mesh = Mesh(np.array(devices).reshape(2, 2), axis_names=('x', 'y'))
# sharding = NamedSharding(mesh, PartitionSpec('x', 'y'))
#
# arr = jax.device_put(jnp.ones((8, 8)), sharding)
#
# # Render the sharding visualization
# fig = treescope.render_array_sharding(arr)
# treescope.display(fig)
#
# # With explicit axis ordering
# fig = treescope.render_array_sharding(
#     arr,
#     rows=[0],     # Show axis 0 as rows
#     columns=[1],  # Show axis 1 as columns
# )
# treescope.display(fig)
```

--------------------------------

### Show Multiple Objects with Treescope

Source: https://context7.com/google-deepmind/treescope/llms.txt

Displays multiple objects inline with rich HTML output, similar to Python's print. Supports options for spacing, word wrapping, and forcing autovisualization.

```python
import treescope
import numpy as np

treescope.basic_interactive_setup()

# Show multiple values inline like print()
treescope.show("Array shape:", np.random.randn(3, 4), "Config:", {"lr": 0.01})

# Show without spaces between items
treescope.show("a", "b", "c", space_separated=False)

# Enable word wrapping at line ends
treescope.show("Long text", np.ones(100), wrap=True)

# Force array autovisualization for this call only
treescope.show(
    "Weights:",
    np.random.randn(32, 32),
    autovisualize=True  # Enable autovisualizer for this call
)
```

--------------------------------

### Implement Custom List Wrapper

Source: https://context7.com/google-deepmind/treescope/llms.txt

Define a class with __treescope_repr__ to customize how sequence-like objects are rendered, including optional styling parameters.

```python
class LayerStack:
    def __init__(self, layers):
        self._layers = list(layers)

    def __getitem__(self, idx):
        return self._layers[idx]

    def __treescope_repr__(self, path, subtree_renderer):
        return repr_lib.render_list_wrapper(
            object_type=type(self),
            wrapped_sequence=self._layers,
            path=path,
            subtree_renderer=subtree_renderer,
            color="#fff3e6",  # Light orange background
        )

stack = LayerStack([NeuralLayer(64, 32), NeuralLayer(32, 16)])
treescope.display(stack)
```

--------------------------------

### Verbose metadata autovisualizer

Source: https://github.com/google-deepmind/treescope/blob/main/docs/notebooks/building_custom_visualizations.ipynb

An autovisualizer that displays additional metadata and path information alongside the array.

```python
def my_verbose_autovisualizer(
    value: Any,
    path: tuple[Any, ...] | None,
):
  if isinstance(value, np.ndarray):
    size = value.size
    token_groups = [
        (id(value) // div) % 1000
        for div in (1000000000000, 1000000, 1000, 1)
    ]
    return treescope.IPythonVisualization(
        treescope.figures.inline(
            "Hello world!\n",
            treescope.render_array(value),
            f"\nThis array contains {size} elements and has Python id {id(value):,}, which you could tokenize as  ",
            treescope.integer_digitbox(token_groups[0], label=str(token_groups[0])),
            "  ", treescope.integer_digitbox(token_groups[1], label=str(token_groups[1])),
            "  ", treescope.integer_digitbox(token_groups[2], label=str(token_groups[2])),
            "  ", treescope.integer_digitbox(token_groups[3], label=str(token_groups[3])),
            f"\nThe path to this node is {path}",
        ),
        replace=False
    )
```

```python
with treescope.active_autovisualizer.set_scoped(
    my_verbose_autovisualizer
):
  IPython.display.display({
      "foo": np.arange(10)[:, None] * np.arange(10)[None, :],
      "bar": np.sin(np.arange(100) * 0.1).reshape((10,10))
  })
```

--------------------------------

### Render Object Constructor Helper

Source: https://github.com/google-deepmind/treescope/blob/main/docs/notebooks/building_custom_visualizations.ipynb

Use `render_object_constructor` for a simple way to create a rendering for your object based on its type and attributes. Ensure `roundtrippable=True` only if the object can be recreated with the provided attributes.

```python
class MySimpleType:
  def __init__(self, foo):
    self.foo = foo

  def __treescope_repr__(self, path, subtree_renderer):
    return treescope.repr_lib.render_object_constructor(
        object_type=type(self),
        attributes={"foo": self.foo},
        path=path,
        subtree_renderer=subtree_renderer,
        # Pass `roundtrippable=True` only if you can rebuild your object by
        # calling `__init__` with these attributes!
        roundtrippable=True,
    )
```

--------------------------------

### Glue Treescope visualization of Penzai model

Source: https://github.com/google-deepmind/treescope/blob/main/docs/_include/_glue_figures.ipynb

Stores the captured HTML output of the Penzai model visualization as 'treescope_penzai' using myst-nb's glue.

```python
myst_nb.glue(
    "treescope_penzai",
    IPython.display.HTML(
        "".join(output.data['text/html'] for output in capturer.outputs)
    ),
)
```

--------------------------------

### IPython and Colab Integration

Source: https://github.com/google-deepmind/treescope/blob/main/docs/notebooks/pretty_printing.ipynb

Configure Treescope as the default formatter or display objects manually in notebook environments.

```python
treescope.register_as_default()
# ^ called by default when running treescope.basic_interactive_setup()
```

```python
treescope.display(["some object"])
```

```python
treescope.show("A value:", ["some object"])
```

```python
["some object"]
```

--------------------------------

### Implement __treescope_ndarray_adapter__ for Array Visualization

Source: https://github.com/google-deepmind/treescope/blob/main/docs/notebooks/building_custom_visualizations.ipynb

Implement this method for array or tensor types to enable automatic visualization. It should return an `NDArrayAdapter` instance for your type.

```python
class MyCustomType:
  ...

  def __treescope_ndarray_adapter__(self) -> NDArrayAdapter:
    ...
```

--------------------------------

### Build Simple Inline Figure with Treescope

Source: https://github.com/google-deepmind/treescope/blob/main/docs/notebooks/building_custom_visualizations.ipynb

Demonstrates building a complex inline figure using `treescope.figures`. This includes styling text, embedding Python objects and arrays, and integrating plots from other libraries like Plotly.

```python
treescope.figures.inline(
    "This is a simple inline output figure. You can ",
    treescope.figures.bolded("emphasize"),
    " parts of the output, or ",
    treescope.figures.with_color("change their color, ", "red"),
    "or even\n"
    "embed Python objects like ",
    [1, 2, 3],
    " or even array visualizations like\n",
    np.linspace(-10, 10, 20),
    ".\nIt's also possible to indent parts of the input, like",
    treescope.figures.indented(treescope.figures.inline(
        "this. Indents apply to\nnewlines and\nembedded objects too:\n",
        [1, 2, 3, 4, 5]
    )),
    "You can also embed figures from other libraries:\n",
    px.histogram(
        jax.random.uniform(jax.random.key(0), (1000,)),
        width=400, height=200
    ).update_layout(
        margin=dict(l=20, r=20, t=20, b=20)
    ),
    "\nAdditionally, you can insert colored \"digitboxes\", which Trescope uses\n",
    "to render token IDs: ",
    treescope.integer_digitbox(1, label="1"),
    " ",
    treescope.integer_digitbox(2, label="2"),
    " ",
    treescope.integer_digitbox(12345, label="12345"),
    ". And you can add ",
    treescope.figures.text_on_color(
        "text with a ",
        value=0.2, vmax=1.0
    ),
    treescope.figures.text_on_color(
        "colormapped",
        value=-1.0, vmax=1.0
    ),
    treescope.figures.text_on_color(
        " background color,",
        value=0.8, vmax=1.0
    ),
    "\nwhich can be useful for showing token probabilities or similar per-token info."
)
```

--------------------------------

### Manage Contextual Values in Treescope

Source: https://context7.com/google-deepmind/treescope/llms.txt

Demonstrates thread-safe management of Treescope settings using global and scoped values. Shows how to set autovisualizers and abbreviation thresholds, including nested contexts with ExitStack.

```python
import treescope
import contextlib

# Set a value globally
treescope.active_autovisualizer.set_globally(treescope.ArrayAutovisualizer())

# Use scoped value (automatically restored after context)
with treescope.active_autovisualizer.set_scoped(None):
    # Autovisualizer is disabled in this block
    treescope.display({"data": [1, 2, 3]})

# Autovisualizer is restored here

# Get current value
current = treescope.active_autovisualizer.get()

# Multiple scoped values with ExitStack
with contextlib.ExitStack() as stack:
    stack.enter_context(treescope.abbreviation_threshold.set_scoped(2))
    stack.enter_context(treescope.active_autovisualizer.set_scoped(None))
    # Both settings active in this block
    treescope.display({"nested": {"data": [1, 2, 3]}})

# Configure abbreviation threshold
treescope.abbreviation_threshold.set_globally(1)  # Abbreviate at depth 1
treescope.roundtrip_abbreviation_threshold.set_globally(2)  # For roundtrip mode
```

--------------------------------

### Configure Array Colormaps

Source: https://context7.com/google-deepmind/treescope/llms.txt

Define global or scoped colormaps for array visualization using RGB tuples.

```python
import treescope
import numpy as np

# Set custom sequential colormap (for non-centered data)
# Colormap is a list of (R, G, B) tuples, values 0-255
custom_sequential = [
    (255, 255, 255),  # White (low)
    (255, 200, 200),
    (255, 100, 100),
    (200, 0, 0),      # Dark red (high)
]
treescope.default_sequential_colormap.set_globally(custom_sequential)

# Set custom diverging colormap (for zero-centered data)
custom_diverging = [
    (0, 0, 200),      # Blue (negative)
    (100, 100, 255),
    (255, 255, 255),  # White (zero)
    (255, 100, 100),
    (200, 0, 0),      # Red (positive)
]
treescope.default_diverging_colormap.set_globally(custom_diverging)

# Use palettable for professional colormaps
# from palettable.matplotlib import Viridis_20, Inferno_20
# treescope.default_sequential_colormap.set_globally(Viridis_20.colors)

# Scoped colormap override
with treescope.default_diverging_colormap.set_scoped(custom_diverging):
    arr = np.random.randn(10, 10)
    fig = treescope.render_array(arr, around_zero=True)
    treescope.display(fig)
```

--------------------------------

### Custom Object Rendering with Treescope

Source: https://context7.com/google-deepmind/treescope/llms.txt

Implement `__treescope_repr__` with helper functions from `repr_lib` for seamless integration of custom types.

```python
from treescope import repr_lib

class MyCustomType:
    def __init__(self, value):
        self.value = value

    def __treescope_repr__(self):
        return repr_lib.object_repr(self, value=self.value)
```

--------------------------------

### Configure Treescope as Default Renderer

Source: https://github.com/google-deepmind/treescope/blob/main/docs/index.rst

Register Treescope as the default pretty printer for the notebook or enable automatic array visualization.

```python
treescope.register_as_default()
```

```python
treescope.active_autovisualizer.set_globally(treescope.ArrayAutovisualizer())
```

```python
treescope.basic_interactive_setup(autovisualize_arrays=True)
```

--------------------------------

### Display Penzai model with Treescope

Source: https://github.com/google-deepmind/treescope/blob/main/docs/_include/_glue_figures.ipynb

Captures the output of Treescope's display for the Penzai model within a specific autovisualizer and expansion strategy.

```python
with IPython.utils.capture.capture_output() as capturer:
  with treescope.active_autovisualizer.set_scoped(treescope.ArrayAutovisualizer()):
    with treescope.using_expansion_strategy(max_height=30):
      treescope.display(model)
```

--------------------------------

### Visualize Built-in Types and Literals

Source: https://github.com/google-deepmind/treescope/blob/main/docs/notebooks/pretty_printing.ipynb

Displays standard Python lists, primitives, and multiline strings.

```python
[
    [1, 2, 3, 4],
    ["a", "b", "c", "d"],
    [True, False, None, NotImplemented, Ellipsis],
    ["a\n  multiline\n    string"]
]
```

--------------------------------

### Continuous mode autovisualizer

Source: https://github.com/google-deepmind/treescope/blob/main/docs/notebooks/building_custom_visualizations.ipynb

An autovisualizer implementation that forces continuous mode rendering for numpy arrays.

```python
def my_continuous_autovisualizer(
    value: Any,
    path: tuple[Any, ...] | None,
):
  if isinstance(value, np.ndarray):
    return treescope.IPythonVisualization(
        treescope.render_array(value, continuous=True, around_zero=False),
        replace=True,
    )
```

```python
with treescope.active_autovisualizer.set_scoped(
    my_continuous_autovisualizer
):
  IPython.display.display({
      "foo": np.arange(10)[:, None] * np.arange(10)[None, :],
      "bar": np.sin(np.arange(100) * 0.1).reshape((10,10))
  })
```

--------------------------------

### Enable autovisualization for a single display call

Source: https://github.com/google-deepmind/treescope/blob/main/docs/api/treescope.rst

Passes the autovisualize argument to display or show functions to enable visualization for a specific call.

```python
treescope.display(..., autovisualize=True)
```

--------------------------------

### Autovisualize Sharding Object

Source: https://github.com/google-deepmind/treescope/blob/main/docs/notebooks/array_visualization.ipynb

Automatically visualize a sharding object itself using the `%%autovisualize` magic command.

```python
%%autovisualize
treescope.display(pos_sharding)
```

--------------------------------

### Glue Treescope visualization of nested parameters

Source: https://github.com/google-deepmind/treescope/blob/main/docs/_include/_glue_figures.ipynb

Stores the captured HTML output of the nested parameters visualization as 'treescope_after' using myst-nb's glue.

```python
myst_nb.glue(
    "treescope_after",
    IPython.display.HTML(
        "".join(output.data['text/html'] for output in capturer.outputs)
    ),
)
```

--------------------------------

### Implement Custom Dictionary Wrapper

Source: https://context7.com/google-deepmind/treescope/llms.txt

Define a class with __treescope_repr__ to customize how dictionary-like objects are rendered in Treescope.

```python
class Config:
    def __init__(self, **kwargs):
        self._data = kwargs

    def __getitem__(self, key):
        return self._data[key]

    def __treescope_repr__(self, path, subtree_renderer):
        return repr_lib.render_dictionary_wrapper(
            object_type=type(self),
            wrapped_dict=self._data,
            path=path,
            subtree_renderer=subtree_renderer,
        )

config = Config(lr=0.001, epochs=100, batch_size=32)
treescope.display(config)
```

--------------------------------

### Render JAX Array Sharding

Source: https://github.com/google-deepmind/treescope/blob/main/docs/notebooks/array_visualization.ipynb

Visualize the sharding of a JAX array using `treescope.render_array_sharding`. This shows how array elements are distributed across devices.

```python
treescope.render_array_sharding(sharded_array)
```

--------------------------------

### Implement __treescope_repr__ for Custom Rendering

Source: https://github.com/google-deepmind/treescope/blob/main/docs/notebooks/building_custom_visualizations.ipynb

Implement this method to define how your custom type is displayed in Treescope. It receives the node's path and a subtree renderer function.

```python
class MyCustomType:
  ...

  def __treescope_repr__(
    self,
    path: str,
    subtree_renderer: Callable[
      [Any, str | None], treescope.rendering_parts.Rendering
    ],
  ) -> treescope.rendering_parts.Rendering | type(NotImplemented):
    ...
```

--------------------------------

### Render negative integers

Source: https://github.com/google-deepmind/treescope/blob/main/docs/notebooks/array_visualization.ipynb

Displays negative integers with a sign indicator.

```python
treescope.render_array(np.arange(21 * 21).reshape((21, 21)) - 220)
```

--------------------------------

### Standard Python Object Rendering

Source: https://github.com/google-deepmind/treescope/blob/main/docs/notebooks/pretty_printing.ipynb

Treescope supports rendering common Python collection types.

```python
[
    [(), (1,), (1, 2, 3)],
    {"foo": "bar", "baz": "qux"},
    {(1,2,3):(4,5,6), (7,8,9):(10,11,12)},
    {"a", "b", "c", "d"}
]
```

--------------------------------

### Offline HTML Rendering

Source: https://github.com/google-deepmind/treescope/blob/main/docs/notebooks/pretty_printing.ipynb

Export Treescope output to static HTML files for offline viewing.

```python
with treescope.active_autovisualizer.set_scoped(treescope.ArrayAutovisualizer()):
  contents = treescope.render_to_html(some_arrays)

with open("/tmp/treescope_output.html", "w") as f:
  f.write(contents)

# Uncomment to download the file:
# import google.colab.files
# google.colab.files.download("/tmp/treescope_output.html")
```

--------------------------------

### Custom Visualization Configuration

Source: https://github.com/google-deepmind/treescope/blob/main/docs/notebooks/pretty_printing.ipynb

Define custom visualization functions to render objects using rich display formats like Plotly.

```python
# You can use most rich display objects, for instance a plotly figure:
import plotly.io
import plotly.express as px

# Treescope uses similar embed settings as Colab, so configure it to render
# like in colab:
plotly.io.renderers.default = "colab"

def visualize_with_histograms(value, path):
  if isinstance(value, (np.ndarray, jax.Array)):
    return treescope.IPythonVisualization(
        px.histogram(
            value.flatten(),
            width=400, height=200
        ).update_layout(
            margin=dict(l=20, r=20, t=20, b=20)
        )
    )
```

```python
with treescope.active_autovisualizer.set_scoped(visualize_with_histograms):
  treescope.display(some_arrays)
```

--------------------------------

### Visualize Functions

Source: https://github.com/google-deepmind/treescope/blob/main/docs/notebooks/pretty_printing.ipynb

Displays various function types including built-ins, lambdas, and JAX-transformed functions.

```python
[
    jnp.sum,
    dataclasses.dataclass,
    lambda x: x + 2,
    jax.vmap(lambda x: x),
]
```

--------------------------------

### Visualize Multidimensional Arrays and Tensors

Source: https://github.com/google-deepmind/treescope/blob/main/docs/notebooks/pretty_printing.ipynb

Summarizes shape, statistics, and contents for JAX, Numpy, and PyTorch arrays, including named axes.

```python
[
    jnp.arange(1000),
    np.array([[np.nan] * 100, [0] * 50 + [1] * 50]),
    torch.linspace(-10, 20, 50),
]
```

```python
from penzai import pz
pz.nx.wrap(
    jax.random.normal(jax.random.key(1), (10, 4, 16)),
).tag("query_seq", "heads", "embed")
```

--------------------------------

### Render Objects to Text

Source: https://context7.com/google-deepmind/treescope/llms.txt

Generates a plain text string representation of an object using Treescope's formatting.

```python
import treescope
import numpy as np

data = {"a": [1, 2, 3], "b": np.array([4, 5, 6])}

# Render to plain text
text = treescope.render_to_text(data)
print(text)
# Output: {'a': [1, 2, 3], 'b': array([4, 5, 6], dtype=int64)}
```

--------------------------------

### Register Treescope as Default Pretty Printer

Source: https://github.com/google-deepmind/treescope/blob/main/README.md

Configure Treescope to be the default pretty printer for your IPython notebook.

```python
treescope.register_as_default()
```

--------------------------------

### Register Autovisualize Magic

Source: https://github.com/google-deepmind/treescope/blob/main/docs/notebooks/array_visualization.ipynb

Enables Treescope's automatic visualization magic commands in the current environment.

```python
import treescope
treescope.register_autovisualize_magic()
```

--------------------------------

### Use autovisualize magic in a cell

Source: https://github.com/google-deepmind/treescope/blob/main/docs/api/treescope.rst

Enables automatic visualization for the duration of a single IPython cell.

```python
%%autovisualize treescope.ArrayAutovisualizer()
treescope.display(...)
```

```python
%%autovisualize
# ^ with no arguments, uses the default array autovisualizer
treescope.display(...)
```

--------------------------------

### Autovisualize Sharded Array

Source: https://github.com/google-deepmind/treescope/blob/main/docs/notebooks/array_visualization.ipynb

Use the `%%autovisualize` magic command to automatically render a sharded array when it's encountered during pretty-printing.

```python
%%autovisualize
treescope.display(sharded_array)
```

--------------------------------

### Register IPython Cell Magics

Source: https://context7.com/google-deepmind/treescope/llms.txt

Enable autovisualization and context management features within Jupyter notebook cells.

```python
import treescope

# Register the magics
treescope.register_autovisualize_magic()
treescope.register_context_manager_magic()

# %%autovisualize magic - enables array autovisualization for the cell
# %%autovisualize
# import numpy as np
# weights = np.random.randn(32, 64)
# weights  # Automatically visualized with ArrayAutovisualizer

# %%autovisualize with explicit autovisualizer
# %%autovisualize treescope.ArrayAutovisualizer(maximum_size=1000)
# large_array = np.random.randn(100, 100)
# large_array

# %%autovisualize False - disable autovisualization
# %%autovisualize False
# array = np.random.randn(10, 10)
# array  # Rendered without visualization

# %%with magic - run cell under a context manager
# %%with treescope.abbreviation_threshold.set_scoped(2)
# nested = {"a": {"b": {"c": {"d": [1, 2, 3]}}}}
# nested  # Abbreviated at depth 2

# %%with treescope.using_expansion_strategy(max_height=5)
# large_structure = {"items": list(range(100))}
# large_structure  # Constrained to 5 lines height
```

--------------------------------

### Render individual integers

Source: https://github.com/google-deepmind/treescope/blob/main/docs/notebooks/array_visualization.ipynb

Displays single integers using digit boxes.

```python
treescope.integer_digitbox(42, size="30px")
```

```python
treescope.integer_digitbox(1234, size="30px")
```

```python
treescope.integer_digitbox(7654321, size="30px")
```

--------------------------------

### Render array with outlier clipping

Source: https://github.com/google-deepmind/treescope/blob/main/docs/notebooks/array_visualization.ipynb

Demonstrates automatic outlier clipping in array visualization.

```python
my_outlier_array = np.cos(np.arange(300).reshape((10,30)) * 0.2)
my_outlier_array[4, 2] = 10.0
treescope.render_array(my_outlier_array)
```

--------------------------------

### Visualize Arbitrary PyTree Types

Source: https://github.com/google-deepmind/treescope/blob/main/docs/notebooks/pretty_printing.ipynb

Uses fallback rendering for JAX-registered PyTree types.

```python
jax.tree_util.Partial(lambda x, y, z: x + y, 10, y=100)
```

--------------------------------

### Apply custom colormaps

Source: https://github.com/google-deepmind/treescope/blob/main/docs/notebooks/array_visualization.ipynb

Uses external color palettes to define the rendering colors.

```python
import palettable
treescope.render_array(my_array, colormap=palettable.matplotlib.Inferno_20.colors)
```

```python
treescope.render_array(my_array, colormap=palettable.cmocean.sequential.Speed_20.colors)
```

--------------------------------

### Rendering to Strings

Source: https://github.com/google-deepmind/treescope/blob/main/docs/api/treescope.rst

Functions to render objects to HTML or text strings for later use.

```APIDOC
## Rendering to Strings

### Description
Convert objects into string representations instead of displaying them directly.

### Functions
- **render_to_html(value)** - Renders the value to an HTML string.
- **render_to_text(value)** - Renders the value to a plain text string.
```

--------------------------------

### Display Objects with Roundtrip Mode

Source: https://github.com/google-deepmind/treescope/blob/main/docs/notebooks/pretty_printing.ipynb

Use roundtrip mode to ensure object representations include qualified names and syntax for potential reconstruction.

```python
my_struct = jax.ShapeDtypeStruct(shape=(20, 10), dtype=jnp.float32)
my_struct
```

```python
treescope.display(my_struct, roundtrip_mode=True)
```

--------------------------------

### Set global autovisualizer

Source: https://github.com/google-deepmind/treescope/blob/main/docs/api/treescope.rst

Configures an autovisualizer to be used for all Treescope outputs.

```python
treescope.active_autovisualizer.set_globally(
    treescope.ArrayAutovisualizer()  # or your own autovisualizer
)
```

--------------------------------

### Render Integer Digitbox

Source: https://context7.com/google-deepmind/treescope/llms.txt

Displays integers as nested colored squares for visual representation.

```python
import treescope

# Render a single integer as digitbox
fig = treescope.integer_digitbox(12345)
treescope.display(fig)

# With custom label
fig = treescope.integer_digitbox(42, label="answer")
treescope.display(fig)

# With custom size
fig = treescope.integer_digitbox(999, size="2em")
treescope.display(fig)
```

--------------------------------

### Nest parameters for Treescope display

Source: https://github.com/google-deepmind/treescope/blob/main/docs/_include/_glue_figures.ipynb

Organizes model parameters into a nested dictionary structure suitable for Treescope's hierarchical visualization.

```python
nested_params = {}
for param in params:
  label_parts = param.label.split("/")
  current = nested_params
  for part in label_parts[:-1]:
    if part not in current:
      current[part] = {}
    current = current[part]
  current[label_parts[-1]] = param.value.data_array
```

--------------------------------

### Visualize Penzai NamedArrays

Source: https://github.com/google-deepmind/treescope/blob/main/docs/notebooks/array_visualization.ipynb

Renders arrays with named axes using Penzai integration.

```python
from penzai import pz
my_named_array = pz.nx.wrap(
    jax.random.normal(jax.random.key(1), (10, 4, 16)),
).tag("query_seq", "heads", "embed")

treescope.render_array(my_named_array, columns=["embed", "heads"])
```

--------------------------------

### Function Reflection and Canonical Aliases

Source: https://github.com/google-deepmind/treescope/blob/main/docs/notebooks/pretty_printing.ipynb

Treescope identifies and summarizes functions, closures, and constants by locating their definitions.

```python
repr(jax.nn.relu)
```

```python
jax.nn.relu
```

```python
jnp.sum
```

```python
def my_function():
  print("hello world!")
```

```python
my_function
```

--------------------------------

### Visualize NamedTuples

Source: https://github.com/google-deepmind/treescope/blob/main/docs/notebooks/pretty_printing.ipynb

Renders a NamedTuple instance with its field values.

```python
class Foo(typing.NamedTuple):
  a: int
  b: str

Foo(a=1, b="bar")
```

--------------------------------

### Render large integer arrays

Source: https://github.com/google-deepmind/treescope/blob/main/docs/notebooks/array_visualization.ipynb

Visualizes large integer arrays using nested box patterns.

```python
treescope.render_array(np.arange(1000).reshape((10,100)))
```

```python
treescope.render_array(
    jnp.arange(20)[:, None] * jnp.arange(20)[None, :]
)
```

--------------------------------

### Configure colormap scaling

Source: https://github.com/google-deepmind/treescope/blob/main/docs/notebooks/array_visualization.ipynb

Adjusts the colormap behavior using around_zero, vmin, and vmax parameters.

```python
treescope.render_array(my_array, around_zero=False)
```

```python
treescope.render_array(my_array, vmax=0.7)
```

```python
treescope.render_array(my_array, vmin=-0.1, vmax=0.7)
```