haliax 1.4.dev407__tar.gz → 1.4.dev409__tar.gz

{haliax-1.4.dev407 → haliax-1.4.dev409}/.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`
@@ -18,10 +18,10 @@ APIs that don't translate well to named tensors are intentionally omitted here.
 - [ ] `atan2`
 - [ ] `average`
 - [ ] `bartlett`
-- [ ] `bitwise_count`
-- [ ] `bitwise_invert`
-- [ ] `bitwise_left_shift`
-- [ ] `bitwise_right_shift`
+- [x] `bitwise_count`
+- [x] `bitwise_invert`
+- [x] `bitwise_left_shift`
+- [x] `bitwise_right_shift`
 - [ ] `blackman`
 - [ ] `block`
 - [ ] `broadcast_shapes`
@@ -93,23 +93,23 @@ 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`
+- [x] `packbits`
 - [ ] `partition`
 - [ ] `percentile`
 - [ ] `permute_dims`
@@ -152,7 +152,7 @@ APIs that don't translate well to named tensors are intentionally omitted here.
 - [ ] `triu_indices`
 - [ ] `triu_indices_from`
 - [ ] `union1d`
-- [ ] `unpackbits`
+- [x] `unpackbits`
 - [ ] `unravel_index`
 - [ ] `unstack`
 - [ ] `unwrap`

{haliax-1.4.dev407 → haliax-1.4.dev409}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: haliax
-Version: 1.4.dev407
+Version: 1.4.dev409
 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.dev409}/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
 
@@ -164,6 +176,8 @@ These are all more or less directly from JAX's NumPy API.
 ::: haliax.arctan
 ::: haliax.arctanh
 ::: haliax.around
+::: haliax.bitwise_count
+::: haliax.bitwise_invert
 ::: haliax.bitwise_not
 ::: haliax.cbrt
 ::: haliax.ceil
@@ -221,7 +235,9 @@ These are all more or less directly from JAX's NumPy API.
 ::: haliax.add
 ::: haliax.arctan2
 ::: haliax.bitwise_and
+::: haliax.bitwise_left_shift
 ::: haliax.bitwise_or
+::: haliax.bitwise_right_shift
 ::: haliax.bitwise_xor
 ::: haliax.divide
 ::: haliax.divmod
@@ -258,6 +274,8 @@ These are all more or less directly from JAX's NumPy API.
 
 ::: haliax.bincount
 ::: haliax.clip
+::: haliax.packbits
+::: haliax.unpackbits
 ::: haliax.isclose
 ::: haliax.allclose
 ::: haliax.array_equal

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

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

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

@@ -83,6 +83,8 @@ from .ops import (
     unique_counts,
     unique_inverse,
     unique_all,
+    packbits,
+    unpackbits,
     searchsorted,
     bincount,
     where,
@@ -337,6 +339,14 @@ def around(a: A) -> A:
     return wrap_elemwise_unary(jnp.around, a)
 
 
+def bitwise_count(a: A) -> A:
+    return wrap_elemwise_unary(jnp.bitwise_count, a)
+
+
+def bitwise_invert(a: A) -> A:
+    return wrap_elemwise_unary(jnp.bitwise_invert, a)
+
+
 def bitwise_not(a: A) -> A:
     return wrap_elemwise_unary(jnp.bitwise_not, a)
 
@@ -560,6 +570,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 +673,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 +768,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)
@@ -717,6 +826,14 @@ def bitwise_and(x1: NamedOrNumeric, x2: NamedOrNumeric, /) -> NamedOrNumeric:
     return jnp.bitwise_and(x1, x2)  # type: ignore
 
 
+@wrap_elemwise_binary
+def bitwise_left_shift(x1: NamedOrNumeric, x2: NamedOrNumeric, /) -> NamedOrNumeric:
+    """
+    Named version of [jax.numpy.bitwise_left_shift](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.bitwise_left_shift.html)
+    """
+    return jnp.bitwise_left_shift(x1, x2)  # type: ignore
+
+
 @wrap_elemwise_binary
 def bitwise_or(x1: NamedOrNumeric, x2: NamedOrNumeric, /) -> NamedOrNumeric:
     """
@@ -725,6 +842,14 @@ def bitwise_or(x1: NamedOrNumeric, x2: NamedOrNumeric, /) -> NamedOrNumeric:
     return jnp.bitwise_or(x1, x2)  # type: ignore
 
 
+@wrap_elemwise_binary
+def bitwise_right_shift(x1: NamedOrNumeric, x2: NamedOrNumeric, /) -> NamedOrNumeric:
+    """
+    Named version of [jax.numpy.bitwise_right_shift](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.bitwise_right_shift.html)
+    """
+    return jnp.bitwise_right_shift(x1, x2)  # type: ignore
+
+
 @wrap_elemwise_binary
 def bitwise_xor(x1: NamedOrNumeric, x2: NamedOrNumeric, /) -> NamedOrNumeric:
     """
@@ -981,6 +1106,8 @@ __all__ = [
     "arctan",
     "arctanh",
     "around",
+    "bitwise_count",
+    "bitwise_invert",
     "bitwise_not",
     "cbrt",
     "ceil",
@@ -1034,12 +1161,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 +1185,8 @@ __all__ = [
     "var",
     "cumsum",
     "cumprod",
+    "nancumprod",
+    "nancumsum",
     "sort",
     "scan",
     "fold",
@@ -1060,6 +1199,8 @@ __all__ = [
     "unique_counts",
     "unique_inverse",
     "unique_all",
+    "packbits",
+    "unpackbits",
     "searchsorted",
     "bincount",
     "clip",
@@ -1068,7 +1209,9 @@ __all__ = [
     "add",
     "arctan2",
     "bitwise_and",
+    "bitwise_left_shift",
     "bitwise_or",
+    "bitwise_right_shift",
     "bitwise_xor",
     "divide",
     "divmod",

{haliax-1.4.dev407 → haliax-1.4.dev409}/src/haliax/ops.py

@@ -508,6 +508,42 @@ def bincount(
     return NamedArray(result, (Counts,))
 
 
+def packbits(a: NamedArray, axis: AxisSelector, *, bitorder: str = "big") -> NamedArray:
+    """Named version of `jax.numpy.packbits`."""
+
+    axis_index = a.axis_indices(axis)
+    if not isinstance(axis_index, int):
+        raise ValueError("packbits only supports a single existing axis")
+
+    result = jnp.packbits(a.array, axis=axis_index, bitorder=bitorder)
+    old_axis = a.axes[axis_index]
+    new_size = (old_axis.size + 7) // 8
+    new_axis = old_axis.resize(new_size)
+    new_axes = a.axes[:axis_index] + (new_axis,) + a.axes[axis_index + 1 :]
+    return NamedArray(result, new_axes)
+
+
+def unpackbits(
+    a: NamedArray,
+    axis: AxisSelector,
+    *,
+    count: int | None = None,
+    bitorder: str = "big",
+) -> NamedArray:
+    """Named version of `jax.numpy.unpackbits`."""
+
+    axis_index = a.axis_indices(axis)
+    if not isinstance(axis_index, int):
+        raise ValueError("unpackbits only supports a single existing axis")
+
+    result = jnp.unpackbits(a.array, axis=axis_index, count=count, bitorder=bitorder)
+    old_axis = a.axes[axis_index]
+    new_size = count if count is not None else old_axis.size * 8
+    new_axis = old_axis.resize(new_size)
+    new_axes = a.axes[:axis_index] + (new_axis,) + a.axes[axis_index + 1 :]
+    return NamedArray(result, new_axes)
+
+
 __all__ = [
     "trace",
     "where",
@@ -517,6 +553,8 @@ __all__ = [
     "pad_left",
     "pad",
     "clip",
+    "packbits",
+    "unpackbits",
     "unique",
     "unique_values",
     "unique_counts",

haliax-1.4.dev409/tests/test_bitwise_ops.py

@@ -0,0 +1,45 @@
+import jax.numpy as jnp
+import haliax as hax
+from haliax import Axis
+
+
+def test_bitwise_count_invert():
+    A = Axis("A", 4)
+    x = hax.named(jnp.array([0, 1, 2, 3], dtype=jnp.uint8), (A,))
+
+    inv = hax.bitwise_invert(x)
+    assert jnp.all(inv.array == jnp.bitwise_invert(x.array))
+
+    cnt = hax.bitwise_count(x)
+    assert jnp.all(cnt.array == jnp.bitwise_count(x.array))
+
+
+def test_bitwise_shift():
+    A = Axis("A", 4)
+    x = hax.named(jnp.array([1, 2, 3, 4], dtype=jnp.int32), (A,))
+    shift = hax.named(jnp.array([1, 1, 1, 1], dtype=jnp.int32), (A,))
+
+    left = hax.bitwise_left_shift(x, shift)
+    assert jnp.all(left.array == jnp.bitwise_left_shift(x.array, shift.array))
+
+    right = hax.bitwise_right_shift(left, shift)
+    assert jnp.all(right.array == jnp.bitwise_right_shift(left.array, shift.array))
+
+    right_scalar = hax.bitwise_right_shift(x, 1)
+    assert jnp.all(right_scalar.array == jnp.bitwise_right_shift(x.array, 1))
+
+
+def test_packbits_unpackbits():
+    B = Axis("B", 10)
+    data = hax.named(jnp.array([0, 1, 1, 0, 1, 0, 0, 1, 1, 0], dtype=jnp.uint8), (B,))
+
+    packed = hax.packbits(data, B)
+    expected_packed = jnp.packbits(data.array, axis=0)
+    assert jnp.all(packed.array == expected_packed)
+    assert packed.axes[0].name == B.name
+    assert packed.axes[0].size == 2
+
+    unpacked = hax.unpackbits(packed, "B", count=B.size)
+    expected_unpacked = jnp.unpackbits(packed.array, axis=0, count=B.size)
+    assert jnp.all(unpacked.array == expected_unpacked)
+    assert unpacked.axes[0].size == B.size

haliax-1.4.dev409/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"