haliax 1.4.dev403__tar.gz → 1.4.dev405__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: haliax
-Version: 1.4.dev403
+Version: 1.4.dev405
 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.dev403 → haliax-1.4.dev405}/docs/api.md

@@ -260,6 +260,7 @@ These are all more or less directly from JAX's NumPy API.
 ::: haliax.clip
 ::: haliax.isclose
 ::: haliax.pad
+::: haliax.searchsorted
 ::: haliax.top_k
 ::: haliax.trace
 ::: haliax.tril

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

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

{haliax-1.4.dev403 → haliax-1.4.dev405}/src/haliax/__init__.py

@@ -80,6 +80,7 @@ from .ops import (
     unique_counts,
     unique_inverse,
     unique_all,
+    searchsorted,
     bincount,
     where,
 )
@@ -1056,6 +1057,7 @@ __all__ = [
     "unique_counts",
     "unique_inverse",
     "unique_all",
+    "searchsorted",
     "bincount",
     "clip",
     "tril",

{haliax-1.4.dev403 → haliax-1.4.dev405}/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.dev405}/src/haliax/ops.py

@@ -430,6 +430,34 @@ def unique_all(
     return values, indices, inverse, counts
 
 
+def searchsorted(
+    a: NamedArray,
+    v: NamedArray | ArrayLike,
+    *,
+    side: str = "left",
+    sorter: NamedArray | ArrayLike | None = None,
+    method: str = "scan",
+) -> NamedArray:
+    """Named version of `jax.numpy.searchsorted`.
+
+    ``a`` and ``sorter`` (if provided) must be one-dimensional.
+    The returned array has the same axes as ``v``.
+    """
+
+    if a.ndim != 1:
+        raise ValueError("searchsorted only supports 1D 'a'")
+
+    if not isinstance(v, NamedArray):
+        v = haliax.named(v, ())
+
+    sorter_arr = None
+    if sorter is not None:
+        sorter_arr = sorter.array if isinstance(sorter, NamedArray) else jnp.asarray(sorter)
+
+    result = jnp.searchsorted(a.array, v.array, side=side, sorter=sorter_arr, method=method)
+    return NamedArray(result, v.axes)
+
+
 def bincount(
     x: NamedArray,
     Counts: Axis,
@@ -471,5 +499,6 @@ __all__ = [
     "unique_counts",
     "unique_inverse",
     "unique_all",
+    "searchsorted",
     "bincount",
 ]

{haliax-1.4.dev403 → haliax-1.4.dev405}/tests/core_test.py

@@ -238,6 +238,26 @@ def test_cumsum_etc():
     assert hax.argsort(named1, axis=Width).axes == (Height, Width, Depth)
 
 
+def test_searchsorted():
+    A = hax.Axis("a", 5)
+    V = hax.Axis("v", 4)
+
+    a = hax.named([1, 3, 5, 7, 9], axis=A)
+    v = hax.named([0, 3, 6, 10], axis=V)
+
+    result = hax.searchsorted(a, v)
+    assert jnp.all(result.array == jnp.searchsorted(a.array, v.array))
+    assert result.axes == (V,)
+
+    unsorted = hax.named([5, 1, 3, 7, 4], axis=A)
+    sorter = hax.argsort(unsorted, axis=A)
+    result = hax.searchsorted(unsorted, v, sorter=sorter, side="right")
+    assert jnp.all(
+        result.array == jnp.searchsorted(unsorted.array, v.array, sorter=sorter.array, side="right")
+    )
+    assert result.axes == (V,)
+
+
 def test_rearrange():
     H, W, D, C = hax.make_axes(H=2, W=3, D=4, C=5)
 

{haliax-1.4.dev403 → haliax-1.4.dev405}/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"