haliax 1.4.dev408__tar.gz → 1.4.dev409__tar.gz

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

@@ -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`
@@ -109,7 +109,7 @@ APIs that don't translate well to named tensors are intentionally omitted here.
 - [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.dev408 → haliax-1.4.dev409}/PKG-INFO

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

@@ -176,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
@@ -233,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
@@ -270,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.dev408 → 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)
 
@@ -816,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:
     """
@@ -824,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:
     """
@@ -1080,6 +1106,8 @@ __all__ = [
     "arctan",
     "arctanh",
     "around",
+    "bitwise_count",
+    "bitwise_invert",
     "bitwise_not",
     "cbrt",
     "ceil",
@@ -1171,6 +1199,8 @@ __all__ = [
     "unique_counts",
     "unique_inverse",
     "unique_all",
+    "packbits",
+    "unpackbits",
     "searchsorted",
     "bincount",
     "clip",
@@ -1179,7 +1209,9 @@ __all__ = [
     "add",
     "arctan2",
     "bitwise_and",
+    "bitwise_left_shift",
     "bitwise_or",
+    "bitwise_right_shift",
     "bitwise_xor",
     "divide",
     "divmod",

{haliax-1.4.dev408 → 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.dev408/src/haliax/__about__.py

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