haliax 1.4.dev407__tar.gz → 1.4.dev408__tar.gz

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

@@ -5,7 +5,7 @@ APIs that don't translate well to named tensors are intentionally omitted here.
 
 ## numpy
 - [x] `allclose`
-- [ ] `amin`
+- [x] `amin`
 - [ ] `append`
 - [ ] `apply_along_axis`
 - [ ] `apply_over_axes`
@@ -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.dev407 → haliax-1.4.dev408}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: haliax
-Version: 1.4.dev407
+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.dev407 → 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
 

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

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

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

@@ -560,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)
@@ -656,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
 
 
@@ -673,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)
@@ -1034,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",
@@ -1048,6 +1157,8 @@ __all__ = [
     "var",
     "cumsum",
     "cumprod",
+    "nancumprod",
+    "nancumsum",
     "sort",
     "scan",
     "fold",

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.dev407/src/haliax/__about__.py

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