treescope_ipython_setup_nested_structure_jax_visualization.py

1import treescope
2import jax
3import jax.numpy as jnp
4
5# 1. Set up Treescope as the default IPython renderer
6treescope.basic_interactive_setup()
7
8# 2. Define a nested data structure (e.g., a JAX model or state)
9data = {
10    "params": {
11        "weights": jnp.ones((4, 4)),
12        "bias": jnp.zeros((4,)),
13    },
14    "metadata": "Example structure",
15    "mask": jnp.array([True, False, True, True]),
16}
17
18# 3. Simply evaluate the object to see the interactive visualization
19# (In a notebook, this will use Treescope's rich display by default)
20data