repository·main·Indexed 19 days ago
https://github.com/google-deepmind/treescopeAn interactive HTML pretty-printer and N-dimensional array visualizer for machine learning research in IPython notebooks. It provides enhanced visibility into neural network models and tensors, supporting NumPy, JAX, PyTorch, and Penzai arrays. Key features include automatic array visualization via ArrayAutovisualizer, interactive tree expansion/collapsing, roundtrip mode for type identification, and specialized rendering for JAX array shardings.
r key to enable roundtrip mode. This mode adds qualified names to every type in the visualization, making it easier to identify the specific types within your object structure.Treescope can automatically trigger specific visualizations for certain leaves in a tree using an "autovisualizer". The default is the ArrayAutovisualizer.
Globally: Set a global autovisualizer for all Treescope outputs:
treescope.active_autovisualizer.set_globally(treescope.ArrayAutovisualizer())Per Call:
Pass the autovisualize argument to a display function:
treescope.display(obj, autovisualize=True)Using IPython Magics:
Use the %%autovisualize magic in a cell. If no argument is provided, it uses the default array autovisualizer:
%%autovisualize
# Uses default array autovisualizer
treescope.display(obj)
# Or specify a custom one
%%autovisualize treescope.ArrayAutovisualizer()
treescope.display(obj)# Global setup
treescope.active_autovisualizer.set_globally(treescope.ArrayAutovisualizer())
# Single call setup
treescope.display(my_array, autovisualize=True)For a quick start that both sets Treescope as the default pretty printer and enables automatic array visualization, use treescope.basic_interactive_setup(autovisualize_arrays=True).
import treescope
treescope.basic_interactive_setup(autovisualize_arrays=True)To configure Treescope as the default renderer for all IPython cell outputs, use treescope.basic_interactive_setup().
This setup performs several actions:
%%autovisualize and %%with IPython magics.For granular control, you can register these components individually using:
treescope.register_as_default()treescope.register_autovisualize_magic()treescope.register_context_manager_magic()import treescope
treescope.basic_interactive_setup()Install Treescope using pip to enable interactive HTML pretty-printing and N-dimensional array visualization in IPython notebooks.
pip install treescopeTreescope provides a set of builder functions to construct a tree of RenderableTreePart objects. These parts can be rendered to text or interactive HTML. You can use these functions within node handlers or the __treescope_repr__ method to define how custom types are visualized.
Note: Do not use the internal RenderableTreePart class directly to build parts; instead, use the exposed builder functions provided in treescope.rendering_parts or higher-level wrappers in treescope.repr_lib.
The ContextualValue[T] class manages a global value that can be read anywhere and modified either globally or within a delimited scope. Local (scoped) values always take precedence over global values. This is useful for reducing boilerplate when you need to change a setting or configuration only within a specific part of your code without affecting the rest of the application.
from treescope.context import ContextualValue
# Initialize with a default value
contextual = ContextualValue(10)
# Access the global value
print(contextual.get()) # Output: 10
# Use a scoped context to temporarily change the value
with contextual.set_scoped(3):
print(contextual.get()) # Output: 3
# Outside the scope, it reverts to the global value
print(contextual.get()) # Output: 10An autovisualizer is a function that Treescope calls on every subtree during rendering. This allows you to automatically inject rich visualizations (like plots or custom array views) into specific types of data within a larger PyTree.
An autovisualizer function must follow this signature:
def autovisualizer_fn(value: Any, path: tuple[Any, ...] | None) -> pz.ts.IPythonVisualization | pz.ts.ChildAutovisualizer | None:
...Return Values:
pz.ts.IPythonVisualization(figure, replace=True/False): Replaces the subtree with a visualization (replace=True) or adds the visualization alongside the subtree (replace=False).pz.ts.ChildAutovisualizer: Tells Treescope to use a different autovisualizer for the children of this node.None: Processes the subtree using the default rendering logic.Usage Patterns:
with treescope.active_autovisualizer.set_scoped(my_autovisualizer): to apply the visualizer to a specific block of code.%%autovisualize magic command in a notebook cell to apply an autovisualizer to the entire cell's output: %%autovisualize my_autovisualizer.def my_continuous_autovisualizer(value, path):
if isinstance(value, np.ndarray):
return treescope.IPythonVisualization(
treescope.render_array(value, continuous=True, around_zero=False),
replace=True,
)
with treescope.active_autovisualizer.set_scoped(my_continuous_autovisualizer):
import IPython
IPython.display.display(np.arange(10))Treescope's output is usually valid Python syntax, but it may lack information needed to rebuild certain objects (like custom types or dataclasses with complex __init__ methods).
Roundtrip mode fixes this by:
< and >).pz.dataclass_from_attributes) to bypass custom __init__ logic.How to enable:
r key.roundtrip_mode=True to the display function.treescope.display(my_struct, roundtrip_mode=True)Most rendering options in Treescope are of type context.ContextualValue. You can manage these in two ways:
context.ContextualValue.set_scoped to apply settings to a specific scope.context.ContextualValue.set_globally to apply settings to all subsequent rendering.Available configuration keys include:
active_rendereractive_autovisualizeractive_expansion_strategydefault_diverging_colormapdefault_sequential_colormapdefault_magic_autovisualizerabbreviation_thresholdroundtrip_abbreviation_thresholdTreescope uses AxisInfo to describe how dimensions in an array can be accessed. When implementing get_axis_info_for_array_data, you can use one of three types of axis information:
PositionalAxisInfo: An ordinary axis accessed by its logical index and size. Common for standard NDArrays.NamedPositionlessAxisInfo: An axis that can only be accessed by a name (e.g., used by penzai.core.named_axes).NamedPositionalAxisInfo: An axis that can be accessed by both its logical index and its name (e.g., used by PyTorch).Note that the axis_logical_index allows supporting 'views' where the axis ordering in the data differs from the logical ordering.
from treescope.ndarray_adapters import PositionalAxisInfo, NamedPositionlessAxisInfo, NamedPositionalAxisInfo
# Example: A positional axis at logical index 0 with size 10
axis1 = PositionalAxisInfo(axis_logical_index=0, size=10)
# Example: A named axis 'batch' with size 32
axis2 = NamedPositionlessAxisInfo(axis_name='batch', size=32)
# Example: A named axis 'channel' at logical index 1 with size 3
axis3 = NamedPositionalAxisInfo(axis_logical_index=1, axis_name='channel', size=3)active_renderer determines the set of handlers and postprocessors used when rendering an object to HTML. Users can override this to change how nodes are rendered. Library functions can retrieve the current value to allow user-configurable rendering, often using TreescopeRenderer.extend_with to add functionality.