haliax 1.4.dev403__tar.gz → 1.4.dev404__tar.gz

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

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

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