haliax 1.4.dev402__tar.gz → 1.4.dev404__tar.gz

{haliax-1.4.dev402 → haliax-1.4.dev404}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: haliax
-Version: 1.4.dev402
+Version: 1.4.dev404
 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.dev404/src/haliax/__about__.py

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

{haliax-1.4.dev402 → haliax-1.4.dev404}/src/haliax/_src/scan.py

@@ -1,9 +1,11 @@
 import dataclasses
 import functools as ft
+import inspect
 from typing import Any, Callable, Literal, ParamSpec, Protocol, Sequence, Tuple, TypeVar, Union, overload
 
 import equinox as eqx
 import jax
+import jax.tree_util as jtu
 from jaxtyping import PyTree
 
 import haliax
@@ -357,7 +359,46 @@ def scan(
             return carry, y
 
         true_axis = _infer_axis_size_from_tree(axis_first_xs, axis)
-        axis_size = _infer_axis_size_from_tree(axis_first_xs, axis).size
+        axis_size = true_axis.size
+
+        # build a mapping from positional argument indices to their names for friendlier error messages
+        sig = inspect.signature(f)
+        arg_pos_names: dict[int, str] = {}
+        params = list(sig.parameters.values())[1:]  # skip carry
+        pos_count = 0
+        var_pos_name: str | None = None
+        for param in params:
+            if param.kind in (
+                inspect.Parameter.POSITIONAL_ONLY,
+                inspect.Parameter.POSITIONAL_OR_KEYWORD,
+            ):
+                arg_pos_names[pos_count] = param.name
+                pos_count += 1
+            elif param.kind == inspect.Parameter.VAR_POSITIONAL:
+                var_pos_name = param.name
+                break
+        if var_pos_name is not None:
+            for i in range(pos_count, len(args)):
+                arg_pos_names[i] = f"{var_pos_name}[{i - pos_count}]"
+
+        path_leaves, _ = jtu.tree_flatten_with_path(axis_first_xs, is_leaf=is_named_array)
+        mismatched = []
+        for path, leaf in path_leaves:
+            if isinstance(leaf, NamedArray):
+                lead_size = leaf.array.shape[0]
+            elif is_jax_array_like(leaf):
+                lead_size = leaf.shape[0]
+            else:
+                continue
+            if lead_size != axis_size:
+                mismatched.append((path, lead_size))
+        if mismatched:
+            details = ", ".join(
+                f"{_format_tree_path(p, arg_pos_names)} has leading dimension {s}" for p, s in mismatched
+            )
+            raise ValueError(
+                f"scan got `length` argument of {axis_size} but some inputs had different leading axis sizes: {details}"
+            )
 
         nested_scan = checkpoint.nested
         outer_block_size = nested_scan_outer_block(nested_scan, axis_size)
@@ -508,6 +549,36 @@ def _zero_if_array_else_none(x: Any) -> ResolvedUnnamedAxisSpec:
     return 0 if is_jax_array_like(x) else None
 
 
+def _format_tree_path(
+    path: tuple[jtu.KeyEntry, ...], arg_pos_names: dict[int, str] | None = None
+) -> str:
+    parts: list[str] = []
+    i = 0
+    if len(path) >= 2 and isinstance(path[0], jtu.SequenceKey):
+        if path[0].idx == 0 and isinstance(path[1], jtu.SequenceKey):
+            name = (arg_pos_names or {}).get(path[1].idx)
+            if name is not None:
+                parts.append(name)
+            else:
+                parts.append(f"[{path[1].idx}]")
+            i = 2
+        elif path[0].idx == 1 and isinstance(path[1], jtu.DictKey):
+            parts.append(str(path[1].key))
+            i = 2
+    for p in path[i:]:
+        if isinstance(p, jtu.GetAttrKey):
+            parts.append("." + p.name)
+        elif isinstance(p, jtu.DictKey):
+            parts.append(f"[{p.key!r}]")
+        elif isinstance(p, jtu.SequenceKey):
+            parts.append(f"[{p.idx}]")
+        else:  # pragma: no cover - future-proofing
+            parts.append(str(p))
+    if parts and parts[0].startswith("."):
+        parts[0] = parts[0][1:]
+    return "".join(parts) or "<root>"
+
+
 def _infer_axis_size_from_tree(result, axis):
     if isinstance(axis, str):
         result_leaves = jax.tree_util.tree_leaves(result, is_leaf=_is_passive_array)

{haliax-1.4.dev402 → haliax-1.4.dev404}/src/haliax/partitioning.py

@@ -1,4 +1,5 @@
 import contextlib
+import dataclasses
 import functools
 import threading
 import typing
@@ -210,6 +211,37 @@ def pspec_for(
                 return None
             else:
                 return pspec_for_axis(node.axes, resource_mapping)
+        elif isinstance(node, eqx.Module):
+            # handle eqx.Module explicitly so that we can look at axis_names metadata
+            updates: dict[str, typing.Any] = {}
+            for field in dataclasses.fields(node):
+                if field.metadata.get("static", False):
+                    continue
+
+                value = getattr(node, field.name)
+                axis_names = field.metadata.get("axis_names") if field.metadata is not None else None
+                if axis_names is not None and is_jax_array_like(value):
+                    current_sharding = (
+                        getattr(value, "sharding", None) if preserve_existing_shardings else None
+                    )
+                    if current_sharding is not None:
+                        updates[field.name] = None
+                    else:
+                        updates[field.name] = pspec_for_axis(axis_names, resource_mapping)
+                else:
+                    updates[field.name] = htu.tree_map(
+                        partition_spec, value, is_leaf=lambda x: isinstance(x, eqx.Module)
+                    )
+
+            new_node = object.__new__(type(node))
+            for field in dataclasses.fields(node):
+                object.__setattr__(
+                    new_node,
+                    field.name,
+                    updates.get(field.name, getattr(node, field.name)),
+                )
+
+            return new_node
         elif is_jax_array_like(node):
             sharding = getattr(node, "sharding", None)
             # TODO: these are usually replicated. Is there a better way to tell?
@@ -232,7 +264,7 @@ def pspec_for(
         else:
             return None
 
-    return htu.tree_map(partition_spec, tree)
+    return htu.tree_map(partition_spec, tree, is_leaf=lambda x: isinstance(x, eqx.Module))
 
 
 def infer_resource_partitions(

{haliax-1.4.dev402 → haliax-1.4.dev404}/tests/test_hof.py

@@ -5,6 +5,7 @@ from jax.random import PRNGKey
 import haliax as hax
 from haliax import Axis, NamedArray
 from haliax.util import is_named_array
+import pytest
 
 
 def test_scan():
@@ -136,6 +137,38 @@ def test_scan_hierarchical():
     assert jnp.all(jnp.equal(selected.array, selected_blocked.array))
 
 
+def test_scan_reports_mismatched_unnamed_array():
+    Height = Axis("Height", 2)
+
+    def f(c, x, y):
+        return c, x + y
+
+    good = jnp.zeros((Height.size, 3))
+    bad = jnp.zeros((Height.size - 1, 3))
+
+    with pytest.raises(ValueError) as e:
+        hax.scan(f, Height)(0, good, y=bad)
+
+    assert "y has leading dimension" in str(e.value)
+
+
+def test_scan_reports_eqx_module_field_path():
+    Height = Axis("Height", 2)
+
+    class Foo(eqx.Module):
+        my_array: jnp.ndarray
+
+    foo = Foo(jnp.zeros((Height.size - 1, 3)))
+
+    def f(c, foo):
+        return c, foo.my_array
+
+    with pytest.raises(ValueError) as e:
+        hax.scan(f, Height)(0, foo)
+
+    assert "foo.my_array" in str(e.value)
+
+
 def test_fold():
     Height = Axis("Height", 10)
     Width = Axis("Width", 3)

{haliax-1.4.dev402 → haliax-1.4.dev404}/tests/test_partitioning.py

@@ -59,6 +59,34 @@ def test_pspec_for_named_axes():
         assert specs.unnamed1 == PartitionSpec(None)
 
 
+class ArrayModule(eqx.Module):
+    arr: Array = hax.field(axis_names=("dim2", "dim3"))
+
+
+def test_pspec_for_plain_array_axis_names():
+    mesh = Mesh(np.array(jax.devices()).reshape(-1, 1), (ResourceAxis.DATA, ResourceAxis.MODEL))
+    with axis_mapping(resource_map), mesh:
+        mod = ArrayModule(jnp.ones((Dim2.size, Dim3.size)))
+
+        specs: ArrayModule = pspec_for(mod, preserve_existing_shardings=False)
+
+        assert specs.arr == PartitionSpec(ResourceAxis.DATA, ResourceAxis.MODEL)
+
+
+class NestedArrayModule(eqx.Module):
+    inner: ArrayModule
+
+
+def test_pspec_for_plain_array_axis_names_nested_module():
+    mesh = Mesh(np.array(jax.devices()).reshape(-1, 1), (ResourceAxis.DATA, ResourceAxis.MODEL))
+    with axis_mapping(resource_map), mesh:
+        mod = NestedArrayModule(ArrayModule(jnp.ones((Dim2.size, Dim3.size))))
+
+        specs: NestedArrayModule = pspec_for(mod, preserve_existing_shardings=False)
+
+        assert specs.inner.arr == PartitionSpec(ResourceAxis.DATA, ResourceAxis.MODEL)
+
+
 class MyModuleInit(eqx.Module):
     named: NamedArray
     unnamed1: Array

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

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