haliax 1.4.dev404__tar.gz → 1.4.dev405__tar.gz

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

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

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