Move arrayviz to the top level of treescope.

This reduces the nesting of arrayviz and removes the need to refer
to arrayviz.arrayviz directly.

PiperOrigin-RevId: 653404499

penzai/pz/ts.py

@@ -16,10 +16,10 @@
 
 # pylint: disable=g-multiple-import,g-importing-member,unused-import
 
-from penzai.treescope.arrayviz.array_autovisualizer import (
+from penzai.treescope.array_autovisualizer import (
     ArrayAutovisualizer,
 )
-from penzai.treescope.arrayviz.arrayviz import (
+from penzai.treescope.arrayviz import (
     default_diverging_colormap,
     default_sequential_colormap,
     integer_digitbox,

penzai/treescope/__init__.py

@@ -30,6 +30,7 @@
 commonly-used treescope functions and classes.
 """
 
+from . import array_autovisualizer
 from . import arrayviz
 from . import autovisualize
 from . import canonical_aliases

penzai/treescope/arrayviz/array_autovisualizer.py → penzai/treescope/array_autovisualizer.py

@@ -23,9 +23,9 @@
 import numpy as np
 from penzai.core import named_axes
 from penzai.core._treescope_handlers import named_axes_handlers
+from penzai.treescope import arrayviz
 from penzai.treescope import autovisualize
 from penzai.treescope import ndarray_summarization
-from penzai.treescope.arrayviz import arrayviz
 from penzai.treescope.foldable_representation import basic_parts
 from penzai.treescope.foldable_representation import common_structures
 from penzai.treescope.foldable_representation import common_styles

penzai/treescope/treescope_ipython.py

@@ -19,12 +19,12 @@
 
 import jax.numpy as jnp
 
+from penzai.treescope import array_autovisualizer
 from penzai.treescope import autovisualize
 from penzai.treescope import context
 from penzai.treescope import default_renderer
 from penzai.treescope import figures
 from penzai.treescope import object_inspection
-from penzai.treescope.arrayviz import array_autovisualizer
 from penzai.treescope.foldable_representation import basic_parts
 from penzai.treescope.foldable_representation import foldable_impl