treescope 0.1.0__py3-none-any.whl

treescope/__init__.py

@@ -0,0 +1,171 @@
+# Copyright 2024 The Treescope Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Treescope: An interactive HTML pretty-printer and array visualizer.
+
+You can configure treescope as the default IPython pretty-printer using ::
+
+  treescope.basic_interactive_setup()
+
+or, for more control: ::
+
+  treescope.register_as_default()
+  treescope.register_autovisualize_magic()
+  treescope.active_autovisualizer.set_globally(
+      treescope.ArrayAutovisualizer()
+  )
+
+You can also pretty-print individual values using `treescope.show` or
+`treescope.display`.
+"""
+from __future__ import annotations
+
+# pylint: disable=g-importing-member,g-multiple-import,unused-import
+
+import typing
+
+from . import _internal
+from . import canonical_aliases
+from . import context
+from . import dataclass_util
+from . import figures
+from . import formatting_util
+from . import handlers
+from . import lowering
+from . import ndarray_adapters
+from . import renderers
+from . import rendering_parts
+from . import repr_lib
+from . import type_registries
+
+from ._internal.api.array_autovisualizer import (
+    ArrayAutovisualizer,
+)
+from ._internal.api.arrayviz import (
+    default_diverging_colormap,
+    default_sequential_colormap,
+    integer_digitbox,
+    render_array,
+    render_array_sharding,
+)
+from ._internal.api.autovisualize import (
+    IPythonVisualization,
+    VisualizationFromTreescopePart,
+    ChildAutovisualizer,
+    Autovisualizer,
+    active_autovisualizer,
+)
+from ._internal.api.default_renderer import (
+    active_renderer,
+    active_expansion_strategy,
+    render_to_html,
+    render_to_text,
+    using_expansion_strategy,
+)
+from ._internal.api.ipython_integration import (
+    default_magic_autovisualizer,
+    basic_interactive_setup,
+    display,
+    register_as_default,
+    register_autovisualize_magic,
+    register_context_manager_magic,
+    show,
+)
+
+if typing.TYPE_CHECKING:
+  Callable = typing.Callable
+
+# Type annotations and docstrings for imported constants. Used by sphinx to
+# generate documentation.
+
+default_diverging_colormap: context.ContextualValue[
+    list[tuple[int, int, int]]
+] = default_diverging_colormap
+"""Default diverging colormap.
+
+Used by `render_array` when ``around_zero`` is True. Intended for user
+customization in an interactive setting.
+"""
+
+default_sequential_colormap: context.ContextualValue[
+    list[tuple[int, int, int]]
+] = default_sequential_colormap
+"""Default sequential colormap.
+
+Used by `render_array` when ``around_zero`` is False. Intended for user
+customization in an interactive setting.
+"""
+
+active_autovisualizer: context.ContextualValue[Autovisualizer | None] = (
+    active_autovisualizer
+)
+"""The active autovisualizer to use when rendering a tree to HTML.
+
+This can be overridden interactively to enable rich visualizations in
+treescope. Users are free to set this to an arbitrary renderer of their
+choice; a common choice is arrayviz's `ArrayAutovisualizer()`.
+"""
+
+active_renderer: context.ContextualValue[renderers.TreescopeRenderer] = (
+    active_renderer
+)
+"""The default renderer to use when rendering a tree to HTML.
+
+This determines the set of handlers and postprocessors to use when
+rendering an object.
+
+This can be overridden locally to reconfigure how nodes are rendered
+with treescope. Users are free to set this to an arbitrary renderer of
+their choice, and programs should not assume the renderer has a
+particular form. Library functions can retrieve the current value of
+this to render objects in a user-configurable way, and can optionally
+make further adjustments using `TreescopeRenderer.extend_with`.
+"""
+
+active_expansion_strategy: context.ContextualValue[
+    Callable[[rendering_parts.RenderableTreePart], None]
+] = active_expansion_strategy
+"""The default expansion strategy to use when rendering a tree to HTML.
+
+Expansion strategies are used to figure out how deeply to unfold an object
+by default. They should operate by setting the expand states of the
+foldable nodes inside the object.
+"""
+
+# Package version.
+__version__ = '0.1.0'
+
+default_magic_autovisualizer: context.ContextualValue[Autovisualizer] = (
+    default_magic_autovisualizer
+)
+"""The default autovisualizer to use for the ``%%autovisualize`` magic.
+
+This can be overridden interactively to customize the autovisualizer
+used by ``%%autovisualize``.
+"""
+
+
+# Set up canonical aliases for the treescope API itself.
+def _setup_canonical_aliases_for_api():
+  import types  # pylint: disable=g-import-not-at-top
+
+  for key, value in globals().items():
+    if isinstance(value, types.FunctionType) or isinstance(value, type):
+      canonical_aliases.add_alias(
+          value, canonical_aliases.ModuleAttributePath(__name__, (key,))
+      )
+
+
+_setup_canonical_aliases_for_api()
+del _setup_canonical_aliases_for_api, typing

treescope/_internal/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2024 The Treescope Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

treescope/_internal/api/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2024 The Treescope Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

treescope/_internal/api/array_autovisualizer.py

@@ -0,0 +1,397 @@
+# Copyright 2024 The Treescope Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""An automatic NDArray visualizer using arrayviz."""
+from __future__ import annotations
+
+import dataclasses
+import sys
+from typing import Any, Callable, Collection
+
+import numpy as np
+from treescope import dtype_util
+from treescope import lowering
+from treescope import ndarray_adapters
+from treescope import rendering_parts
+from treescope import type_registries
+from treescope._internal import arrayviz_impl
+from treescope._internal.api import arrayviz
+from treescope._internal.api import autovisualize
+
+
+PositionalAxisInfo = ndarray_adapters.PositionalAxisInfo
+NamedPositionlessAxisInfo = ndarray_adapters.NamedPositionlessAxisInfo
+NamedPositionalAxisInfo = ndarray_adapters.NamedPositionalAxisInfo
+AxisInfo = ndarray_adapters.AxisInfo
+
+ArrayInRegistry = Any
+
+
+def _supported_dtype(dtype: np.dtype | None):
+  return dtype is not None and (
+      dtype_util.is_integer_dtype(dtype)
+      or dtype_util.is_floating_dtype(dtype)
+      or np.issubdtype(dtype, np.bool_)
+  )
+
+
+@dataclasses.dataclass
+class ArrayAutovisualizer:
+  """An automatic visualizer for arrays.
+
+  ArrayAutovisualizer supports any array type registered with an NDArrayAdapter
+  in the global type registry.
+
+  Attributes:
+    maximum_size: Maximum numer of elements of an array to show. Arrays larger
+      than this will be truncated along one or more axes.
+    cutoff_size_per_axis: Maximum number of elements of each individual axis to
+      show without truncation. Any axis longer than this will be truncated, with
+      their visual size increasing logarithmically with the true axis size
+      beyond this point.
+    edge_items: How many values to keep along each axis for truncated arrays.
+    prefers_column: Names that should always be assigned to columns.
+    prefers_row: Names that should always be assigned to rows.
+    around_zero: Whether to center continous data around zero.
+    force_continuous: Whether to always render integer arrays as continuous.
+    include_repr_line_threshold: A threshold such that, if the `repr` of the
+      array has fewer than that many lines, we will include that repr in the
+      visualization. Useful for seeing small array values.
+    token_lookup_fn: Optional function that looks up token IDs and adds them to
+      the visualization on hover.
+  """
+
+  maximum_size: int = 4_000
+  cutoff_size_per_axis: int = 128
+  edge_items: int = 5
+  prefers_column: Collection[str] = ()
+  prefers_row: Collection[str] = ()
+  around_zero: bool = True
+  force_continuous: bool = False
+  include_repr_line_threshold: int = 5
+  token_lookup_fn: Callable[[int], str] | None = None
+
+  def _autovisualize_array(
+      self,
+      array: ArrayInRegistry,
+      adapter: ndarray_adapters.NDArrayAdapter,
+      path: str | None,
+      label: str,
+      expand_state: rendering_parts.ExpandState,
+  ) -> rendering_parts.RenderableTreePart:
+    """Helper to visualize an array."""
+    # Extract information about axis names, indices, and sizes.
+    array_axis_info = adapter.get_axis_info_for_array_data(array)
+
+    # Assign axes, using preferred axes if possible.
+    row_axes = []
+    column_axes = []
+    for info in array_axis_info:
+      if isinstance(info, NamedPositionalAxisInfo | NamedPositionlessAxisInfo):
+        if info.axis_name in self.prefers_column:
+          column_axes.append(info.axis_name)
+        elif info.axis_name in self.prefers_row:
+          row_axes.append(info.axis_name)
+
+    # Infer a good truncated shape for this array.
+    edge_items_per_axis = arrayviz_impl.infer_balanced_truncation(
+        tuple(info.size for info in array_axis_info),
+        maximum_size=self.maximum_size,
+        cutoff_size_per_axis=self.cutoff_size_per_axis,
+        minimum_edge_items=self.edge_items,
+    )
+
+    row_axes, column_axes = arrayviz_impl.infer_rows_and_columns(
+        all_axes=array_axis_info,
+        known_rows=row_axes,
+        known_columns=column_axes,
+        edge_items_per_axis=edge_items_per_axis,
+    )
+
+    # Obtain truncated array and mask data from the adapter.
+    truncated_array_data, truncated_mask_data = (
+        adapter.get_array_data_with_truncation(
+            array=array,
+            mask=None,
+            edge_items_per_axis=edge_items_per_axis,
+        )
+    )
+
+    # Maybe infer value labels from a tokenizer.
+    if (
+        self.token_lookup_fn
+        and not self.force_continuous
+        and dtype_util.is_integer_dtype(truncated_array_data.dtype)
+    ):
+      tokens = np.unique(truncated_array_data.flatten()).tolist()
+      value_item_labels = {
+          token: self.token_lookup_fn(token) for token in tokens
+      }
+    else:
+      value_item_labels = None
+
+    array_rendering = arrayviz._render_pretruncated(  # pylint: disable=protected-access
+        array_axis_info=array_axis_info,
+        row_infos=row_axes,
+        column_infos=column_axes,
+        slider_infos=(),
+        truncated_array_data=truncated_array_data,
+        truncated_mask_data=truncated_mask_data,
+        edge_items_per_axis=edge_items_per_axis,
+        continuous="auto",
+        around_zero=self.around_zero,
+        vmax=None,
+        vmin=None,
+        trim_outliers=True,
+        dynamic_colormap="auto",
+        colormap=None,
+        axis_item_labels=None,
+        value_item_labels=value_item_labels,
+        axis_labels=None,
+    )
+    outputs = [array_rendering]
+    last_line_parts = []
+
+    # Render the sharding as well.
+    sharding_info = adapter.get_sharding_info_for_array_data(array)
+    if sharding_info:
+      num_devices = len(sharding_info.device_index_to_shard_slices)
+      if num_devices == 1:
+        [device_id] = sharding_info.device_index_to_shard_slices.keys()
+        last_line_parts.append(
+            f"| Device: {sharding_info.device_type} {device_id}"
+        )
+      else:
+        if sharding_info.fully_replicated:
+          sharding_summary_str = (
+              "Replicated across"
+              f" {num_devices} {sharding_info.device_type} devices"
+          )
+        else:
+          sharding_summary_str = (
+              "Sharded across"
+              f" {num_devices} {sharding_info.device_type} devices"
+          )
+        sharding_rendering = arrayviz.render_array_sharding(
+            array,
+            columns=[c.logical_key() for c in column_axes],
+            rows=[r.logical_key() for r in row_axes],
+        ).treescope_part
+        outputs.append(
+            rendering_parts.build_custom_foldable_tree_node(
+                label=rendering_parts.abbreviation_color(
+                    rendering_parts.siblings(
+                        rendering_parts.text(sharding_summary_str),
+                        rendering_parts.fold_condition(
+                            expanded=rendering_parts.text(":"),
+                            collapsed=rendering_parts.text(
+                                " (click to expand)"
+                            ),
+                        ),
+                    )
+                ),
+                contents=rendering_parts.fold_condition(
+                    expanded=rendering_parts.indented_children(
+                        [sharding_rendering]
+                    ),
+                ),
+            )
+        )
+
+    # We render it with a path, but remove the copy path button. This will be
+    # added back by the caller.
+    if last_line_parts:
+      last_line = rendering_parts.siblings(
+          rendering_parts.fold_condition(
+              expanded=rendering_parts.text("".join(last_line_parts)),
+          ),
+          rendering_parts.text(">"),
+      )
+    else:
+      last_line = rendering_parts.text(">")
+    custom_rendering = rendering_parts.build_custom_foldable_tree_node(
+        label=rendering_parts.abbreviation_color(label),
+        contents=rendering_parts.siblings(
+            rendering_parts.fold_condition(
+                expanded=rendering_parts.indented_children(outputs)
+            ),
+            rendering_parts.abbreviation_color(last_line),
+        ),
+        path=path,
+        expand_state=expand_state,
+    )
+    return custom_rendering.renderable
+
+  def __call__(
+      self, value: Any, path: str | None
+  ) -> autovisualize.VisualizationFromTreescopePart | None:
+    """Implementation of an autovisualizer, visualizing arrays."""
+    # Retrieve the adapter for this array, which we will use to construct
+    # the rendering.
+    adapter = type_registries.lookup_ndarray_adapter(value)
+    if adapter is not None:
+      if not adapter.should_autovisualize(value):
+        return None
+
+      # This is an array we can visualize!
+      # Extract information about axis names, indices, and sizes.
+      array_axis_info = adapter.get_axis_info_for_array_data(value)
+      total_size = np.prod([ax.size for ax in array_axis_info])
+      if total_size == 1:
+        # Don't visualize scalars.
+        return None
+
+      np_dtype = adapter.get_numpy_dtype(value)
+      if not _supported_dtype(np_dtype):
+        return None
+
+      def _placeholder() -> rendering_parts.RenderableTreePart:
+        summary = adapter.get_array_summary(value, fast=True)
+        return rendering_parts.deferred_placeholder_style(
+            rendering_parts.text(f"<{summary}>")
+        )
+
+      def _thunk(
+          expand_state: rendering_parts.ExpandState | None,
+      ) -> rendering_parts.RenderableTreePart:
+        # Full rendering of the array.
+        if expand_state is None:
+          expand_state = rendering_parts.ExpandState.WEAKLY_EXPANDED
+        summary = adapter.get_array_summary(value, fast=False)
+        label = rendering_parts.abbreviation_color(
+            rendering_parts.text(f"<{summary}")
+        )
+        return self._autovisualize_array(
+            value, adapter, path, label, expand_state
+        )
+
+      return autovisualize.VisualizationFromTreescopePart(
+          rendering_parts.RenderableAndLineAnnotations(
+              renderable=lowering.maybe_defer_rendering(
+                  _thunk, _placeholder, expanded_newlines_for_layout=8
+              ),
+              annotations=rendering_parts.build_copy_button(path),
+          )
+      )
+
+    # Not an array in the registry. But it might be a JAX sharding that we can
+    # visualize (if JAX is imported).
+    if "jax" in sys.modules:
+      import jax  # pylint: disable=g-import-not-at-top
+
+      if isinstance(
+          value,
+          jax.sharding.PositionalSharding
+          | jax.sharding.NamedSharding
+          | jax.sharding.Mesh,
+      ):
+        raw_repr = repr(value)
+        repr_oneline = " ".join(line.strip() for line in raw_repr.split("\n"))
+
+        if isinstance(value, jax.sharding.PositionalSharding):
+          sharding = value
+          fake_axis_info = [
+              PositionalAxisInfo(i, size)
+              for i, size in enumerate(sharding.shape)
+          ]
+        elif isinstance(value, jax.sharding.NamedSharding):
+          sharding = value
+          # Named shardings still act on positional arrays, so show them for
+          # the positional shape they require.
+          fake_sizes = []
+          for part in value.spec:
+            if part is None:
+              fake_sizes.append(1)
+            elif isinstance(part, str):
+              fake_sizes.append(value.mesh.shape[part])
+            else:
+              size = int(np.prod([value.mesh.shape[a] for a in part]))
+              fake_sizes.append(size)
+          fake_axis_info = [
+              PositionalAxisInfo(i, size) for i, size in enumerate(fake_sizes)
+          ]
+        else:
+          # Meshes are based on named axes. We build a temporary positional
+          # sharding for visualization, but keep track of name order.
+          assert isinstance(value, jax.sharding.Mesh)
+          mesh = value
+          sharding = jax.sharding.NamedSharding(
+              value, jax.sharding.PartitionSpec(*mesh.axis_names)
+          )
+          fake_axis_info = [
+              NamedPositionlessAxisInfo(name, mesh.shape[name])
+              for name in mesh.axis_names
+          ]
+
+        fake_shape = tuple(ax.size for ax in fake_axis_info)
+        some_device = next(iter(sharding.device_set))
+        device_index_map = sharding.devices_indices_map(fake_shape)
+        sharding_info = ndarray_adapters.ShardingInfo(
+            shard_shape=sharding.shard_shape(fake_shape),
+            device_index_to_shard_slices={
+                d.id: v for d, v in device_index_map.items()
+            },
+            device_type=some_device.platform.upper(),
+            fully_replicated=sharding.is_fully_replicated,
+        )
+        shardvis = arrayviz.render_sharding_info(
+            array_axis_info=fake_axis_info,
+            sharding_info=sharding_info,
+        ).treescope_part
+        custom_rendering = rendering_parts.build_custom_foldable_tree_node(
+            label=rendering_parts.abbreviation_color(
+                rendering_parts.text("<" + repr_oneline)
+            ),
+            contents=rendering_parts.siblings(
+                rendering_parts.fold_condition(
+                    expanded=rendering_parts.indented_children([shardvis])
+                ),
+                rendering_parts.abbreviation_color(rendering_parts.text(">")),
+            ),
+            path=path,
+            expand_state=rendering_parts.ExpandState.EXPANDED,
+        )
+        return autovisualize.VisualizationFromTreescopePart(custom_rendering)
+      else:
+        return None
+
+  @classmethod
+  def for_tokenizer(cls, tokenizer: Any):
+    """Builds an autovisualizer for a tokenizer.
+
+    This method constructs an ArrayAutovisualizer that annotates integer array
+    elements with their token strings. This can then be used to autovisualize
+    tokenized arrays.
+
+    Args:
+      tokenizer: A tokenizer to use. Either a callable mapping token IDs to
+        strings, or a SentencePieceProcessor.
+
+    Returns:
+      An ArrayAutovisualizer that annotates integer array elements with their
+      token strings.
+    """
+    if callable(tokenizer):
+      return cls(token_lookup_fn=lambda x: repr(tokenizer(x)))
+    elif hasattr(tokenizer, "IdToPiece") and hasattr(tokenizer, "GetPieceSize"):
+
+      def lookup(x):
+        if x >= 0 and x < tokenizer.GetPieceSize():
+          return repr(tokenizer.IdToPiece(x))
+        else:
+          return f"<out of bounds: {x}>"
+
+      return cls(token_lookup_fn=lookup)
+    else:
+      raise ValueError(f"Unknown tokenizer type: {tokenizer}")