haliax 1.4.dev410__tar.gz → 1.4.dev411__tar.gz

{haliax-1.4.dev410 → haliax-1.4.dev411}/.agents/projects/api_parity.md

@@ -107,7 +107,7 @@ APIs that don't translate well to named tensors are intentionally omitted here.
 - [x] `nanstd`
 - [x] `nansum`
 - [x] `nanvar`
-- [ ] `nonzero`
+- [x] `nonzero`
 - [ ] `ogrid`
 - [x] `packbits`
 - [ ] `partition`

{haliax-1.4.dev410 → haliax-1.4.dev411}/PKG-INFO

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

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

{haliax-1.4.dev410 → haliax-1.4.dev411}/src/haliax/__about__.py

@@ -3,4 +3,4 @@
 # SPDX-License-Identifier: Apache-2.0
 
 
-__version__ = "1.4.dev410"
+__version__ = "1.4.dev411"

{haliax-1.4.dev410 → haliax-1.4.dev411}/src/haliax/__init__.py

@@ -84,6 +84,7 @@ from .ops import (
     trace,
     tril,
     triu,
+    nonzero,
     unique,
     unique_values,
     unique_counts,
@@ -1199,6 +1200,7 @@ __all__ = [
     "map",
     "vmap",
     "trace",
+    "nonzero",
     "where",
     "unique",
     "unique_values",

{haliax-1.4.dev410 → haliax-1.4.dev411}/src/haliax/ops.py

@@ -108,6 +108,30 @@ def where(
     return NamedArray(raw, condition.axes)
 
 
+def nonzero(array: NamedArray, *, size: Axis, fill_value: int = 0) -> tuple[NamedArray, ...]:
+    """Like :func:`jax.numpy.nonzero`, but with named axes.
+
+    Args:
+        array: The input array to test for nonzero values. Must be a :class:`NamedArray`.
+        size: Axis specifying the size of the output axis. This is required because
+            JAX requires the size of the result at tracing time.
+        fill_value: Value used to fill the output when fewer than ``size`` elements are
+            nonzero. Defaults to ``0``.
+
+    Returns:
+        A tuple of :class:`NamedArray` objects, one for each axis of ``array``. Each
+        returned array has ``size`` as its only axis and contains the indices of the
+        nonzero elements along the corresponding input axis.
+    """
+
+    if not isinstance(array, NamedArray):
+        raise ValueError("array must be a NamedArray")
+
+    return tuple(
+        NamedArray(idx, (size,)) for idx in jnp.nonzero(array.array, size=size.size, fill_value=fill_value)
+    )
+
+
 def clip(array: NamedOrNumeric, a_min: NamedOrNumeric, a_max: NamedOrNumeric) -> NamedArray:
     """Like jnp.clip, but with named axes. This version currently only accepts the three argument form."""
     (array, a_min, a_max), axes = broadcast_arrays_and_return_axes(array, a_min, a_max)

{haliax-1.4.dev410 → haliax-1.4.dev411}/tests/test_ops.py

@@ -161,6 +161,29 @@ def test_where(use_jit):
     assert jnp.all(unnamed_8 == named8.array)
 
 
+@pytest.mark.parametrize("use_jit", [False, True])
+def test_nonzero(use_jit):
+    Height = Axis("Height", 3)
+    Width = Axis("Width", 4)
+
+    data = jnp.array([[0, 1, 0, 2], [3, 0, 4, 0], [0, 0, 5, 6]])
+    named = hax.named(data, (Height, Width))
+
+    Nonzero = Axis("Nonzero", Height.size * Width.size)
+
+    hax_nonzero: Callable = hax.nonzero
+    if use_jit:
+        hax_nonzero = hax.named_jit(hax_nonzero)
+
+    i, j = hax_nonzero(named, size=Nonzero, fill_value=-1)
+    ui, uj = jnp.nonzero(data, size=Nonzero.size, fill_value=-1)
+
+    assert i.axes == (Nonzero,)
+    assert j.axes == (Nonzero,)
+    assert jnp.all(i.array == ui)
+    assert jnp.all(j.array == uj)
+
+
 def test_clip():
     Height = Axis("Height", 10)
     Width = Axis("Width", 3)