haliax 1.4.dev388__tar.gz → 1.4.dev390__tar.gz

{haliax-1.4.dev388 → haliax-1.4.dev390}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: haliax
-Version: 1.4.dev388
+Version: 1.4.dev390
 Summary: Named Tensors for Legible Deep Learning in JAX
 Project-URL: Homepage, https://github.com/stanford-crfm/haliax
 Project-URL: Bug Tracker, https://github.com/stanford-crfm/haliax/issues/

{haliax-1.4.dev388 → haliax-1.4.dev390}/docs/faq.md

@@ -9,3 +9,9 @@ Currently, we diagnose:
 
 * Reuse of arrays or NamedArrays in a field. [Equinox modules must be trees.](https://docs.kidger.site/equinox/faq/#a-module-saved-in-two-places-has-become-two-independent-copies)
 * Use of arrays or NamedArrays in a static field. Static data in JAX/Equinox must be hashable, and arrays are not hashable.
+
+## Tip 2: `hax.debug.visualize_shardings`
+
+Use `hax.debug.visualize_shardings` to quickly inspect how a PyTree is sharded.
+It prints the sharding of each array leaf, including the mapping from named axes
+to physical axes for :class:`haliax.NamedArray` leaves.

haliax-1.4.dev390/src/haliax/__about__.py

@@ -0,0 +1 @@
+__version__ = "1.4.dev390"

{haliax-1.4.dev388 → haliax-1.4.dev390}/src/haliax/_src/scan.py

@@ -33,7 +33,7 @@ class ScanFn(Protocol[Carry, Args, Y]):
         ...
 
 
-@dataclasses.dataclass
+@dataclasses.dataclass(frozen=True)
 class ScanCheckpointPolicy:
     """
     A class that represents a gradient checkpoint policy for blocks in a Stacked module. This is used to control

{haliax-1.4.dev388 → haliax-1.4.dev390}/src/haliax/debug.py

@@ -1,11 +1,14 @@
 import dataclasses
-from typing import List, Tuple, Union
+from typing import List, Tuple, Union, Sequence
 
 import equinox as eqx
+import jax
 import jax.numpy as jnp
 import jax.tree_util as jtu
 
+
 from haliax.core import NamedArray
+from haliax.axis import Axis
 from haliax.util import is_jax_or_hax_array_like
 
 from ._src.util import IdentityMap
@@ -117,3 +120,70 @@ def _check_for_static_arrays(problems, module):
 
     if static_arrays:
         problems.static_arrays.extend(static_arrays)
+
+
+def _pspec_parts(spec_part) -> str:
+    if spec_part is None:
+        return "unsharded"
+    elif isinstance(spec_part, (tuple, list)):
+        return "+".join(str(p) for p in spec_part)
+    else:
+        return str(spec_part)
+
+
+def visualize_named_sharding(axes: Sequence[Axis], sharding: jax.sharding.Sharding) -> None:
+    """Visualize the sharding for a set of named axes.
+
+    This extends :func:`jax.debug.visualize_sharding` to handle arrays with more
+    than two dimensions by falling back to a textual description when necessary.
+    """
+
+    try:
+        pspec = sharding.spec  # type: ignore[attr-defined]
+    except Exception:
+        pspec = (None,) * len(axes)
+
+    parts = [_pspec_parts(p) for p in pspec]
+    num_sharded = sum(p != "unsharded" for p in parts)
+
+    if num_sharded <= 2:
+        try:
+            jax.debug.visualize_sharding([ax.size for ax in axes], sharding)
+        except Exception:
+            pass
+
+    mapping = ", ".join(f"{ax.name}->{part}" for ax, part in zip(axes, parts))
+    print(mapping)
+
+
+def visualize_shardings(tree) -> None:
+    """Print the sharding for each array-like leaf in ``tree``.
+
+    Both :class:`NamedArray` and regular JAX arrays are supported. NamedArrays
+    will show the mapping from logical axis names to physical axes. Plain arrays
+    will fall back to :func:`jax.debug.visualize_sharding`.
+    """
+
+    import haliax.tree_util as htu
+
+    def _show(x):
+        if isinstance(x, NamedArray):
+            arr = x.array
+            axes = x.axes
+        else:
+            arr = x
+            axes = None
+
+        def cb(sh):
+            if axes is not None:
+                visualize_named_sharding(axes, sh)
+            else:
+                try:
+                    jax.debug.visualize_sharding(arr.shape, sh)
+                except Exception:
+                    pass
+
+        jax.debug.inspect_array_sharding(arr, callback=cb)
+        return x
+
+    htu.tree_map(_show, tree, is_leaf=is_jax_or_hax_array_like)

haliax-1.4.dev390/tests/test_visualize_sharding.py

@@ -0,0 +1,67 @@
+import numpy as np
+import jax
+import jax.numpy as jnp
+
+import haliax as hax
+from haliax import Axis
+from haliax.partitioning import ResourceAxis, axis_mapping, named_jit
+from test_utils import skip_if_not_enough_devices
+from haliax.debug import visualize_shardings
+
+Dim1 = Axis("dim1", 8)
+Dim2 = Axis("dim2", 8)
+Dim3 = Axis("dim3", 2)
+
+resource_map = {
+    "dim1": ResourceAxis.DATA,
+    "dim2": ResourceAxis.MODEL,
+    "dim3": ResourceAxis.REPLICA,
+}
+
+
+def test_visualize_shardings_runs(capsys):
+    mesh = jax.sharding.Mesh(
+        np.array(jax.devices()).reshape(-1, 1, 1),
+        (ResourceAxis.DATA, ResourceAxis.MODEL, ResourceAxis.REPLICA),
+    )
+    with axis_mapping(resource_map), mesh:
+        arr = hax.ones((Dim1, Dim2, Dim3))
+        visualize_shardings(arr)
+
+    out = capsys.readouterr().out
+    assert "dim1" in out and "dim2" in out and "dim3" in out
+
+
+def test_visualize_shardings_inside_jit(capsys):
+    mesh = jax.sharding.Mesh(np.array(jax.devices()).reshape(-1, 1), (ResourceAxis.DATA, ResourceAxis.MODEL))
+
+    @named_jit(out_axis_resources={"dim1": ResourceAxis.DATA})
+    def fn(x):
+        visualize_shardings(x)
+        return x
+
+    with axis_mapping({"dim1": ResourceAxis.DATA}), mesh:
+        x = hax.ones(Dim1)
+        fn(x)
+
+    out = capsys.readouterr().out
+    assert "dim1" in out
+
+
+def test_visualize_shardings_plain_array(capsys):
+    x = jnp.ones((4, 4))
+    visualize_shardings(x)
+    out = capsys.readouterr().out
+    assert out.strip() != ""
+
+
+@skip_if_not_enough_devices(2)
+def test_visualize_shardings_model_axis(capsys):
+    devices = jax.devices()
+    mesh = jax.sharding.Mesh(np.array(devices).reshape(-1, 2), (ResourceAxis.DATA, ResourceAxis.MODEL))
+    with axis_mapping({"dim1": ResourceAxis.DATA, "dim2": ResourceAxis.MODEL}), mesh:
+        arr = hax.ones((Dim1, Dim2))
+        visualize_shardings(arr)
+
+    out = capsys.readouterr().out
+    assert "dim2" in out

haliax-1.4.dev388/src/haliax/__about__.py

@@ -1 +0,0 @@
-__version__ = "1.4.dev388"