haliax 1.4.dev406__tar.gz → 1.4.dev408__tar.gz

{haliax-1.4.dev406 → haliax-1.4.dev408}/.agents/projects/api_parity.md

@@ -4,15 +4,15 @@ This document tracks JAX NumPy functions not yet wrapped by Haliax.
 APIs that don't translate well to named tensors are intentionally omitted here. This includes dtype constructors, raw array converters (e.g. `from_dlpack`), indexing helpers like `c_`/`r_`, and functions whose JAX counterparts already work with `NamedArray` out of the box. Basic array construction is handled by `haliax.named`, so functions such as `array` or `asarray` are not listed.
 
 ## numpy
-- [ ] `allclose`
-- [ ] `amin`
+- [x] `allclose`
+- [x] `amin`
 - [ ] `append`
 - [ ] `apply_along_axis`
 - [ ] `apply_over_axes`
 - [ ] `argpartition`
 - [ ] `argwhere`
-- [ ] `array_equal`
-- [ ] `array_equiv`
+- [x] `array_equal`
+- [x] `array_equiv`
 - [ ] `array_split`
 - [ ] `astype`
 - [ ] `atan2`
@@ -93,20 +93,20 @@ APIs that don't translate well to named tensors are intentionally omitted here.
 - [ ] `modf`
 - [ ] `moveaxis`
 - [ ] `nan_to_num`
-- [ ] `nanargmax`
-- [ ] `nanargmin`
-- [ ] `nancumprod`
-- [ ] `nancumsum`
-- [ ] `nanmax`
-- [ ] `nanmean`
+- [x] `nanargmax`
+- [x] `nanargmin`
+- [x] `nancumprod`
+- [x] `nancumsum`
+- [x] `nanmax`
+- [x] `nanmean`
 - [ ] `nanmedian`
-- [ ] `nanmin`
+- [x] `nanmin`
 - [ ] `nanpercentile`
-- [ ] `nanprod`
+- [x] `nanprod`
 - [ ] `nanquantile`
-- [ ] `nanstd`
-- [ ] `nansum`
-- [ ] `nanvar`
+- [x] `nanstd`
+- [x] `nansum`
+- [x] `nanvar`
 - [ ] `nonzero`
 - [ ] `ogrid`
 - [ ] `packbits`

{haliax-1.4.dev406 → haliax-1.4.dev408}/PKG-INFO

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

@@ -128,12 +128,22 @@ You can convert it to a [jax.numpy.ndarray][] with [haliax.NamedArray.scalar][],
 
 ::: haliax.all
 ::: haliax.amax
+::: haliax.amin
 ::: haliax.any
 ::: haliax.argmax
 ::: haliax.argmin
 ::: haliax.max
 ::: haliax.mean
 ::: haliax.min
+::: haliax.nanargmax
+::: haliax.nanargmin
+::: haliax.nanmax
+::: haliax.nanmean
+::: haliax.nanmin
+::: haliax.nanprod
+::: haliax.nanstd
+::: haliax.nansum
+::: haliax.nanvar
 ::: haliax.prod
 ::: haliax.ptp
 ::: haliax.std
@@ -146,6 +156,8 @@ don't reduce it.
 
 ::: haliax.cumsum
 ::: haliax.cumprod
+::: haliax.nancumprod
+::: haliax.nancumsum
 ::: haliax.sort
 ::: haliax.argsort
 
@@ -259,6 +271,9 @@ These are all more or less directly from JAX's NumPy API.
 ::: haliax.bincount
 ::: haliax.clip
 ::: haliax.isclose
+::: haliax.allclose
+::: haliax.array_equal
+::: haliax.array_equiv
 ::: haliax.pad
 ::: haliax.searchsorted
 ::: haliax.top_k

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

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

{haliax-1.4.dev406 → haliax-1.4.dev408}/src/haliax/__init__.py

@@ -69,6 +69,9 @@ from .hof import fold, map, scan, vmap
 from .jax_utils import tree_checkpoint_name
 from .ops import (
     clip,
+    allclose,
+    array_equal,
+    array_equiv,
     isclose,
     pad_left,
     pad,
@@ -557,6 +560,13 @@ def amax(array: NamedArray, axis: Optional[AxisSelection] = None, *, where: Opti
     return wrap_reduction_call(jnp.amax, array, axis, where, single_axis_only=False, supports_where=True)
 
 
+def amin(array: NamedArray, axis: Optional[AxisSelection] = None, *, where: Optional[NamedArray] = None) -> NamedArray:
+    """
+    Aliax for min. See min for details.
+    """
+    return wrap_reduction_call(jnp.amin, array, axis, where, single_axis_only=False, supports_where=True)
+
+
 def any(array: NamedArray, axis: Optional[AxisSelection] = None, *, where: Optional[NamedArray] = None) -> NamedArray:
     """True if any elements along a given axis or axes are True. If axis is None, any elements are True."""
     return wrap_reduction_call(jnp.any, array, axis, where, single_axis_only=False, supports_where=True)
@@ -653,6 +663,84 @@ def var(
     )
 
 
+def nanargmax(array: NamedArray, axis: Optional[AxisSelector] = None) -> NamedArray:
+    return wrap_reduction_call(jnp.nanargmax, array, axis, None, single_axis_only=True, supports_where=False)
+
+
+def nanargmin(array: NamedArray, axis: Optional[AxisSelector] = None) -> NamedArray:
+    return wrap_reduction_call(jnp.nanargmin, array, axis, None, single_axis_only=True, supports_where=False)
+
+
+def nanmax(
+    array: NamedArray,
+    axis: Optional[AxisSelection] = None,
+    *,
+    where: Optional[NamedArray] = None,
+) -> NamedArray:
+    return wrap_reduction_call(jnp.nanmax, array, axis, where, single_axis_only=False, supports_where=True)
+
+
+def nanmean(
+    array: NamedArray,
+    axis: Optional[AxisSelection] = None,
+    *,
+    where: Optional[NamedArray] = None,
+    dtype: Optional[DTypeLike] = None,
+) -> NamedArray:
+    return wrap_reduction_call(jnp.nanmean, array, axis, where, single_axis_only=False, supports_where=True, dtype=dtype)
+
+
+def nanmin(
+    array: NamedArray,
+    axis: Optional[AxisSelection] = None,
+    *,
+    where: Optional[NamedArray] = None,
+) -> NamedArray:
+    return wrap_reduction_call(jnp.nanmin, array, axis, where, single_axis_only=False, supports_where=True)
+
+
+def nanprod(
+    array: NamedArray,
+    axis: Optional[AxisSelection] = None,
+    *,
+    where: Optional[NamedArray] = None,
+    dtype: Optional[DTypeLike] = None,
+) -> NamedArray:
+    return wrap_reduction_call(jnp.nanprod, array, axis, where, single_axis_only=False, supports_where=True, dtype=dtype)
+
+
+def nanstd(
+    array: NamedArray,
+    axis: Optional[AxisSelection] = None,
+    *,
+    where: Optional[NamedArray] = None,
+    ddof: int = 0,
+    dtype: Optional[DTypeLike] = None,
+) -> NamedArray:
+    return wrap_reduction_call(jnp.nanstd, array, axis, where, single_axis_only=False, supports_where=True, dtype=dtype, ddof=ddof)
+
+
+def nansum(
+    array: NamedArray,
+    axis: Optional[AxisSelection] = None,
+    *,
+    where: Optional[NamedArray] = None,
+    dtype: Optional[DTypeLike] = None,
+) -> NamedArray:
+    return wrap_reduction_call(jnp.nansum, array, axis, where, single_axis_only=False, supports_where=True, dtype=dtype)
+
+
+def nanvar(
+    array: NamedArray,
+    axis: Optional[AxisSelection] = None,
+    *,
+    where: Optional[NamedArray] = None,
+    ddof: int = 0,
+    dtype: Optional[DTypeLike] = None,
+) -> NamedArray:
+    return wrap_reduction_call(jnp.nanvar, array, axis, where, single_axis_only=False, supports_where=True, dtype=dtype, ddof=ddof)
+
+
 # "Normalization" functions that use an axis but don't change the shape
 
 
@@ -670,6 +758,20 @@ def cumprod(a: NamedArray, axis: AxisSelector, dtype: Optional[DTypeLike] = None
     return wrap_axiswise_call(jnp.cumprod, a, axis, dtype=dtype, single_axis_only=True)
 
 
+def nancumsum(a: NamedArray, axis: AxisSelector, *, dtype: Optional[DTypeLike] = None) -> NamedArray:
+    """
+    Named version of [jax.numpy.nancumsum](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.nancumsum.html)
+    """
+    return wrap_axiswise_call(jnp.nancumsum, a, axis, dtype=dtype, single_axis_only=True)
+
+
+def nancumprod(a: NamedArray, axis: AxisSelector, dtype: Optional[DTypeLike] = None) -> NamedArray:
+    """
+    Named version of [jax.numpy.nancumprod](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.nancumprod.html)
+    """
+    return wrap_axiswise_call(jnp.nancumprod, a, axis, dtype=dtype, single_axis_only=True)
+
+
 def sort(a: NamedArray, axis: AxisSelector) -> NamedArray:
     """
     Named version of [jax.numpy.sort](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.sort.html)
@@ -1031,12 +1133,22 @@ __all__ = [
     "trunc",
     "all",
     "amax",
+    "amin",
     "any",
     "argmax",
     "argmin",
     "max",
     "mean",
     "min",
+    "nanargmax",
+    "nanargmin",
+    "nanmax",
+    "nanmean",
+    "nanmin",
+    "nanprod",
+    "nanstd",
+    "nansum",
+    "nanvar",
     "prod",
     "product",
     "ptp",
@@ -1045,6 +1157,8 @@ __all__ = [
     "var",
     "cumsum",
     "cumprod",
+    "nancumprod",
+    "nancumsum",
     "sort",
     "scan",
     "fold",
@@ -1105,6 +1219,9 @@ __all__ = [
     "shard",
     "enable_shape_checks",
     "are_shape_checks_enabled",
+    "allclose",
+    "array_equal",
+    "array_equiv",
     "isclose",
     "pad_left",
     "pad",

{haliax-1.4.dev406 → haliax-1.4.dev408}/src/haliax/ops.py

@@ -138,6 +138,29 @@ def isclose(a: NamedArray, b: NamedArray, rtol=1e-05, atol=1e-08, equal_nan=Fals
     return NamedArray(jnp.isclose(a.array, b.array, rtol=rtol, atol=atol, equal_nan=equal_nan), a.axes)
 
 
+def allclose(a: NamedArray, b: NamedArray, rtol=1e-05, atol=1e-08, equal_nan=False) -> bool:
+    """Returns True if two arrays are element-wise equal within a tolerance."""
+    a, b = broadcast_arrays(a, b)
+    return bool(jnp.allclose(a.array, b.array, rtol=rtol, atol=atol, equal_nan=equal_nan))
+
+
+def array_equal(a: NamedArray, b: NamedArray) -> bool:
+    """Returns True if two arrays have the same shape and elements."""
+    if set(a.axes) != set(b.axes):
+        return False
+    b = b.rearrange(a.axes)
+    return bool(jnp.array_equal(a.array, b.array))
+
+
+def array_equiv(a: NamedArray, b: NamedArray) -> bool:
+    """Returns True if two arrays are shape-consistent and equal."""
+    try:
+        a, b = broadcast_arrays(a, b)
+    except ValueError:
+        return False
+    return bool(jnp.array_equal(a.array, b.array))
+
+
 def pad_left(array: NamedArray, axis: Axis, new_axis: Axis, value=0) -> NamedArray:
     """Pad an array along named axes."""
     amount_to_pad_to = new_axis.size - axis.size

haliax-1.4.dev408/tests/test_nan_reductions.py

@@ -0,0 +1,59 @@
+from typing import Any, Callable
+
+import jax.numpy as jnp
+import haliax as hax
+
+
+def _sample_array():
+    Height, Width = hax.make_axes(Height=2, Width=3)
+    data = jnp.array([[1.0, jnp.nan, 3.0], [jnp.nan, 5.0, 6.0]])
+    arr = hax.named(data, (Height, Width))
+    return Height, Width, data, arr
+
+
+def test_amin_alias():
+    Height, Width = hax.make_axes(Height=2, Width=3)
+    data = jnp.arange(6.0).reshape(2, 3)
+    arr = hax.named(data, (Height, Width))
+    assert jnp.array_equal(hax.amin(arr).array, jnp.amin(data))
+    assert jnp.array_equal(hax.amin(arr, axis=Height).array, jnp.amin(data, axis=0))
+    assert hax.amin(arr, axis=Height).axes == (Width,)
+
+
+def test_nan_reductions():
+    Height, Width, data, arr = _sample_array()
+
+    funcs: list[tuple[Callable[..., Any], Callable[..., Any]]] = [
+        (hax.nanmin, jnp.nanmin),
+        (hax.nanmax, jnp.nanmax),
+        (hax.nanmean, jnp.nanmean),
+        (hax.nansum, jnp.nansum),
+        (hax.nanprod, jnp.nanprod),
+        (hax.nanstd, jnp.nanstd),
+        (hax.nanvar, jnp.nanvar),
+    ]
+
+    for hfunc, jfunc in funcs:
+        assert jnp.allclose(hfunc(arr).array, jfunc(data), equal_nan=True)
+        assert jnp.allclose(hfunc(arr, axis=Height).array, jfunc(data, axis=0), equal_nan=True)
+        assert hfunc(arr, axis=Height).axes == (Width,)
+
+    arg_funcs: list[tuple[Callable[..., Any], Callable[..., Any]]] = [
+        (hax.nanargmax, jnp.nanargmax),
+        (hax.nanargmin, jnp.nanargmin),
+    ]
+
+    for hfunc, jfunc in arg_funcs:
+        out = hfunc(arr, axis=Height)
+        assert jnp.array_equal(out.array, jfunc(data, axis=0))
+        assert out.axes == (Width,)
+
+    axiswise_funcs: list[tuple[Callable[..., Any], Callable[..., Any]]] = [
+        (hax.nancumsum, jnp.nancumsum),
+        (hax.nancumprod, jnp.nancumprod),
+    ]
+
+    for hfunc, jfunc in axiswise_funcs:
+        out = hfunc(arr, axis=Height)
+        assert jnp.allclose(out.array, jfunc(data, axis=0), equal_nan=True)
+        assert out.axes == (Height, Width)

{haliax-1.4.dev406 → haliax-1.4.dev408}/tests/test_ops.py

@@ -425,6 +425,32 @@ def test_bincount():
     assert jnp.allclose(out_w.array, expected_w)
 
 
+def test_allclose_array_equal_equiv():
+    A = Axis("A", 2)
+    B = Axis("B", 3)
+    x = hax.random.uniform(PRNGKey(0), (A, B))
+    y = x + 1e-6
+
+    assert hax.allclose(x, y)
+    assert not hax.allclose(x, x + 1.0)
+
+    x1 = hax.ones((A, B))
+    y_reordered = x1.rearrange((B, A))
+    assert hax.array_equal(x1, y_reordered)
+
+    scalar = hax.ones(())
+    assert hax.array_equiv(x1, scalar)
+    assert not hax.array_equal(x1, scalar)
+
+    y_vec = hax.ones((B,))
+    assert hax.array_equiv(x1, y_vec)
+    assert not hax.array_equal(x1, y_vec)
+
+    C = Axis("C", 4)
+    z = hax.ones((C,))
+    assert not hax.array_equiv(x1, z)
+
+
 def test_roll_scalar_named_shift():
     H = Axis("H", 4)
     W = Axis("W", 3)

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

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