haliax 1.4.dev382__tar.gz → 1.4.dev388__tar.gz

haliax-1.4.dev388/.playbooks/wrap-non-named.md

@@ -0,0 +1,49 @@
+# Wrapping Functions with NamedArray Support
+
+This playbook explains how to convert a regular JAX function that works on unnamed arrays into a Haliax function that accepts `NamedArray` inputs and returns `NamedArray` outputs.
+
+## When is wrapping needed?
+Many JAX primitives only operate on regular arrays. To integrate them in Haliax you should provide a thin wrapper that handles axis metadata. Simple elementwise operations and reductions have helper utilities.
+
+## Elemwise Unary
+For a unary function that acts elementwise (e.g. `jnp.abs`):
+
+```python
+from haliax import wrap_elemwise_unary
+
+def abs(a):
+    return wrap_elemwise_unary(jnp.abs, a)
+```
+
+This preserves axis order and dtype.
+
+## Elemwise Binary
+For binary operations (e.g. `jnp.add`), decorate a function with `wrap_elemwise_binary`:
+
+```python
+from haliax import wrap_elemwise_binary
+
+@wrap_elemwise_binary
+def add(x1, x2):
+    return jnp.add(x1, x2)
+```
+
+Broadcasting between `NamedArray`s is handled automatically.
+
+## Reductions
+Reductions require choosing axes to eliminate. Use `wrap_reduction_call`:
+
+```python
+from haliax import wrap_reduction_call
+
+def sum(a, axis=None):
+    return wrap_reduction_call(jnp.sum, a, axis)
+```
+
+`axis` can be an `AxisSelector` or tuple. The wrapper returns a `NamedArray` with those axes removed.
+
+## Harder Cases
+Some functions need bespoke handling. For example `jnp.unique` returns several arrays and may change shape unpredictably. There is no generic helper, so you will need to manually map between `NamedArray` axes and the outputs. Use the lower level utilities in `haliax.wrap` for broadcasting and axis lookup.
+
+## Testing
+Add tests to ensure that named and unnamed calls produce the same results and that axis names are preserved or removed correctly.

{haliax-1.4.dev382 → haliax-1.4.dev388}/AGENTS.md

@@ -17,6 +17,7 @@ repository. Follow these notes when implementing new features or fixing bugs.
 ## Playbook
 
 - Adding Haliax-style tensor typing annotations are described in @.playbooks/add-types.md
+- Wrapping standard JAX functions so they operate on `NamedArray` is explained in @.playbooks/wrap-non-named.md
 
 ## Code Style
 

{haliax-1.4.dev382 → haliax-1.4.dev388}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: haliax
-Version: 1.4.dev382
+Version: 1.4.dev388
 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.dev388/src/haliax/__about__.py

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

{haliax-1.4.dev382 → haliax-1.4.dev388}/src/haliax/__init__.py

@@ -66,7 +66,21 @@ from .core import (
 from .haxtyping import Named
 from .hof import fold, map, scan, vmap
 from .jax_utils import tree_checkpoint_name
-from .ops import clip, isclose, pad_left, pad, trace, tril, triu, where
+from .ops import (
+    clip,
+    isclose,
+    pad_left,
+    pad,
+    trace,
+    tril,
+    triu,
+    unique,
+    unique_values,
+    unique_counts,
+    unique_inverse,
+    unique_all,
+    where,
+)
 from .partitioning import auto_sharded, axis_mapping, fsdp, named_jit, shard, shard_with_axis_mapping
 from .specialized_fns import top_k
 from .types import Scalar
@@ -1034,6 +1048,11 @@ __all__ = [
     "vmap",
     "trace",
     "where",
+    "unique",
+    "unique_values",
+    "unique_counts",
+    "unique_inverse",
+    "unique_all",
     "clip",
     "tril",
     "triu",

{haliax-1.4.dev382 → haliax-1.4.dev388}/src/haliax/ops.py

@@ -3,6 +3,9 @@ from typing import Mapping, Optional, Union
 
 import jax
 import jax.numpy as jnp
+from jaxtyping import ArrayLike
+
+import haliax
 
 from .axis import Axis, AxisSelector, axis_name
 from .core import NamedArray, NamedOrNumeric, broadcast_arrays, broadcast_arrays_and_return_axes, named
@@ -196,4 +199,249 @@ def raw_array_or_scalar(x: NamedOrNumeric):
     return x
 
 
-__all__ = ["trace", "where", "tril", "triu", "isclose", "pad_left", "pad", "clip"]
+@typing.overload
+def unique(
+    array: NamedArray, Unique: Axis, *, axis: AxisSelector | None = None, fill_value: ArrayLike | None = None
+) -> NamedArray:
+    ...
+
+
+@typing.overload
+def unique(
+    array: NamedArray,
+    Unique: Axis,
+    *,
+    return_index: typing.Literal[True],
+    axis: AxisSelector | None = None,
+    fill_value: ArrayLike | None = None,
+) -> tuple[NamedArray, NamedArray]:
+    ...
+
+
+@typing.overload
+def unique(
+    array: NamedArray,
+    Unique: Axis,
+    *,
+    return_inverse: typing.Literal[True],
+    axis: AxisSelector | None = None,
+    fill_value: ArrayLike | None = None,
+) -> tuple[NamedArray, NamedArray]:
+    ...
+
+
+@typing.overload
+def unique(
+    array: NamedArray,
+    Unique: Axis,
+    *,
+    return_counts: typing.Literal[True],
+    axis: AxisSelector | None = None,
+    fill_value: ArrayLike | None = None,
+) -> tuple[NamedArray, NamedArray]:
+    ...
+
+
+@typing.overload
+def unique(
+    array: NamedArray,
+    Unique: Axis,
+    *,
+    return_index: bool = False,
+    return_inverse: bool = False,
+    return_counts: bool = False,
+    axis: AxisSelector | None = None,
+    fill_value: ArrayLike | None = None,
+) -> NamedArray | tuple[NamedArray, ...]:
+    ...
+
+
+def unique(
+    array: NamedArray,
+    Unique: Axis,
+    *,
+    return_index: bool = False,
+    return_inverse: bool = False,
+    return_counts: bool = False,
+    axis: AxisSelector | None = None,
+    fill_value: ArrayLike | None = None,
+) -> NamedArray | tuple[NamedArray, ...]:
+    """
+    Like jnp.unique, but with named axes.
+
+    Args:
+        array: The input array.
+        Unique: The name of the axis that will be created to hold the unique values.
+        fill_value: The value to use for the fill_value argument of jnp.unique
+        axis: The axis along which to find unique values.
+        return_index: If True, return the indices of the unique values.
+        return_inverse: If True, return the indices of the input array that would reconstruct the unique values.
+    """
+    size = Unique.size
+
+    is_multireturn = return_index or return_inverse or return_counts
+
+    kwargs = dict(
+        size=size,
+        fill_value=fill_value,
+        return_index=return_index,
+        return_inverse=return_inverse,
+        return_counts=return_counts,
+    )
+
+    if axis is not None:
+        axis_index = array._lookup_indices(axis)
+        if axis_index is None:
+            raise ValueError(f"Axis {axis} not found in array. Available axes: {array.axes}")
+        out = jnp.unique(array.array, axis=axis_index, **kwargs)
+    else:
+        out = jnp.unique(array.array, **kwargs)
+
+    if is_multireturn:
+        unique = out[0]
+        next_index = 1
+        if return_index:
+            index = out[next_index]
+            next_index += 1
+        if return_inverse:
+            inverse = out[next_index]
+            next_index += 1
+        if return_counts:
+            counts = out[next_index]
+            next_index += 1
+    else:
+        unique = out
+
+    ret = []
+
+    if axis is not None:
+        out_axes = haliax.axis.replace_axis(array.axes, axis, Unique)
+    else:
+        out_axes = (Unique,)
+
+    unique_values = haliax.named(unique, out_axes)
+    if not is_multireturn:
+        return unique_values
+
+    ret.append(unique_values)
+
+    if return_index:
+        ret.append(haliax.named(index, Unique))
+
+    if return_inverse:
+        if axis is not None:
+            assert axis_index is not None
+            inverse = haliax.named(inverse, array.axes[axis_index])
+        else:
+            inverse = haliax.named(inverse, array.axes)
+        ret.append(inverse)
+
+    if return_counts:
+        ret.append(haliax.named(counts, Unique))
+
+    return tuple(ret)
+
+
+def unique_values(
+    array: NamedArray,
+    Unique: Axis,
+    *,
+    axis: AxisSelector | None = None,
+    fill_value: ArrayLike | None = None,
+) -> NamedArray:
+    """Shortcut for :func:`unique` that returns only unique values."""
+
+    return typing.cast(
+        NamedArray,
+        unique(
+            array,
+            Unique,
+            axis=axis,
+            fill_value=fill_value,
+        ),
+    )
+
+
+def unique_counts(
+    array: NamedArray,
+    Unique: Axis,
+    *,
+    axis: AxisSelector | None = None,
+    fill_value: ArrayLike | None = None,
+) -> tuple[NamedArray, NamedArray]:
+    """Shortcut for :func:`unique` that also returns counts."""
+
+    values, counts = typing.cast(
+        tuple[NamedArray, NamedArray],
+        unique(
+            array,
+            Unique,
+            return_counts=True,
+            axis=axis,
+            fill_value=fill_value,
+        ),
+    )
+    return values, counts
+
+
+def unique_inverse(
+    array: NamedArray,
+    Unique: Axis,
+    *,
+    axis: AxisSelector | None = None,
+    fill_value: ArrayLike | None = None,
+) -> tuple[NamedArray, NamedArray]:
+    """Shortcut for :func:`unique` that also returns inverse indices."""
+
+    values, inverse = typing.cast(
+        tuple[NamedArray, NamedArray],
+        unique(
+            array,
+            Unique,
+            return_inverse=True,
+            axis=axis,
+            fill_value=fill_value,
+        ),
+    )
+    return values, inverse
+
+
+def unique_all(
+    array: NamedArray,
+    Unique: Axis,
+    *,
+    axis: AxisSelector | None = None,
+    fill_value: ArrayLike | None = None,
+) -> tuple[NamedArray, NamedArray, NamedArray, NamedArray]:
+    """Shortcut for :func:`unique` returning values, indices, inverse, and counts."""
+
+    values, indices, inverse, counts = typing.cast(
+        tuple[NamedArray, NamedArray, NamedArray, NamedArray],
+        unique(
+            array,
+            Unique,
+            return_index=True,
+            return_inverse=True,
+            return_counts=True,
+            axis=axis,
+            fill_value=fill_value,
+        ),
+    )
+    return values, indices, inverse, counts
+
+
+__all__ = [
+    "trace",
+    "where",
+    "tril",
+    "triu",
+    "isclose",
+    "pad_left",
+    "pad",
+    "clip",
+    "unique",
+    "unique_values",
+    "unique_counts",
+    "unique_inverse",
+    "unique_all",
+]

{haliax-1.4.dev382 → haliax-1.4.dev388}/tests/test_ops.py

@@ -1,4 +1,5 @@
 from typing import Callable
+import typing
 
 import jax.numpy as jnp
 import pytest
@@ -250,3 +251,159 @@ def test_pad():
     assert padded.axes[0].size == Height.size + 3
     assert padded.axes[1].size == Width.size + 1
     assert jnp.all(expected == padded.array)
+
+
+def test_unique():
+    # named version of this test:
+    # >>> M = jnp.array([[1, 2],
+    # ...                [2, 3],
+    # ...                [1, 2]])
+    # >>> jnp.unique(M)
+    # Array([1, 2, 3], dtype=int32)
+
+    Height = Axis("Height", 3)
+    Width = Axis("Width", 2)
+
+    named1 = hax.named([[1, 2], [2, 3], [1, 2]], (Height, Width))
+
+    U = Axis("U", 3)
+
+    named2 = hax.unique(named1, U)
+
+    assert jnp.all(jnp.equal(named2.array, jnp.array([1, 2, 3])))
+
+    #     If you pass an ``axis`` keyword, you can find unique *slices* of the array along
+    #     that axis:
+    #
+    #     >>> jnp.unique(M, axis=0)
+    #     Array([[1, 2],
+    #            [2, 3]], dtype=int32)
+
+    U2 = Axis("U2", 2)
+    named3 = hax.unique(named1, U2, axis=Height)
+    assert jnp.all(jnp.equal(named3.array, jnp.array([[1, 2], [2, 3]])))
+
+    #     >>> x = jnp.array([3, 4, 1, 3, 1])
+    #     >>> values, indices = jnp.unique(x, return_index=True)
+    #     >>> print(values)
+    #     [1 3 4]
+    #     >>> print(indices)
+    #     [2 0 1]
+    #     >>> jnp.all(values == x[indices])
+    #     Array(True, dtype=bool)
+
+    x = hax.named([3, 4, 1, 3, 1], ("Height",))
+    U3 = Axis("U3", 3)
+    values, indices = hax.unique(x, U3, return_index=True)
+
+    assert jnp.all(jnp.equal(values.array, jnp.array([1, 3, 4])))
+    assert jnp.all(jnp.equal(indices.array, jnp.array([2, 0, 1])))
+
+    assert jnp.all(jnp.equal(values.array, x[{"Height": indices}].array))
+
+    #   If you set ``return_inverse=True``, then ``unique`` returns the indices within the
+    #     unique values for every entry in the input array:
+    #
+    #     >>> x = jnp.array([3, 4, 1, 3, 1])
+    #     >>> values, inverse = jnp.unique(x, return_inverse=True)
+    #     >>> print(values)
+    #     [1 3 4]
+    #     >>> print(inverse)
+    #     [1 2 0 1 0]
+    #     >>> jnp.all(values[inverse] == x)
+    #     Array(True, dtype=bool)
+
+    values, inverse = hax.unique(x, U3, return_inverse=True)
+
+    assert jnp.all(jnp.equal(values.array, jnp.array([1, 3, 4])))
+    assert jnp.all(jnp.equal(inverse.array, jnp.array([1, 2, 0, 1, 0])))
+
+    #     In multiple dimensions, the input can be reconstructed using
+    #     :func:`jax.numpy.take`:
+    #
+    #     >>> values, inverse = jnp.unique(M, axis=0, return_inverse=True)
+    #     >>> jnp.all(jnp.take(values, inverse, axis=0) == M)
+    #     Array(True, dtype=bool)
+    #
+
+    values, inverse = hax.unique(named1, U3, axis=Height, return_inverse=True)
+
+    assert jnp.all((values[{"U3": inverse}] == named1).array)
+
+    #     **Returning counts**
+    #  If you set ``return_counts=True``, then ``unique`` returns the number of occurrences
+    #     within the input for every unique value:
+    #
+    #     >>> x = jnp.array([3, 4, 1, 3, 1])
+    #     >>> values, counts = jnp.unique(x, return_counts=True)
+    #     >>> print(values)
+    #     [1 3 4]
+    #     >>> print(counts)
+    #     [2 2 1]
+    #
+    #     For multi-dimensional arrays, this also returns a 1D array of counts
+    #     indicating number of occurrences along the specified axis:
+    #
+    #     >>> values, counts = jnp.unique(M, axis=0, return_counts=True)
+    #     >>> print(values)
+    #     [[1 2]
+    #      [2 3]]
+    #     >>> print(counts)
+    #     [2 1]
+
+    values, counts = hax.unique(x, U3, return_counts=True)
+
+    assert jnp.all(jnp.equal(values.array, jnp.array([1, 3, 4])))
+    assert jnp.all(jnp.equal(counts.array, jnp.array([2, 2, 1])))
+
+    values, counts = hax.unique(named1, U2, axis=Height, return_counts=True)
+
+    assert jnp.all(jnp.equal(values.array, jnp.array([[1, 2], [2, 3]])))
+    assert jnp.all(jnp.equal(counts.array, jnp.array([2, 1])))
+
+
+def test_unique_shortcuts():
+    Height = Axis("Height", 3)
+    Width = Axis("Width", 2)
+
+    arr2d = hax.named([[1, 2], [2, 3], [1, 2]], (Height, Width))
+    U = Axis("U", 3)
+
+    # unique_values
+    uv = hax.unique_values(arr2d, U)
+    uv_expected = hax.unique(arr2d, U)
+    assert jnp.all(uv.array == uv_expected.array)
+
+    # unique_counts
+    vc, cc = hax.unique_counts(arr2d, U)
+    vc_exp, cc_exp = hax.unique(arr2d, U, return_counts=True)
+    assert jnp.all(vc.array == vc_exp.array)
+    assert jnp.all(cc.array == cc_exp.array)
+
+    # unique_inverse
+    Height1 = Axis("Height1", 5)
+    arr1d = hax.named([3, 4, 1, 3, 1], (Height1,))
+    U2 = Axis("U2", 3)
+    vi, ii = hax.unique_inverse(arr1d, U2)
+    vi_exp, ii_exp = hax.unique(arr1d, U2, return_inverse=True)
+    assert jnp.all(vi.array == vi_exp.array)
+    assert jnp.all(ii.array == ii_exp.array)
+
+    # unique_all
+    U3 = Axis("U3", 2)
+    va, ia, ina, ca = hax.unique_all(arr2d, U3, axis=Height)
+    va_exp, ia_exp, ina_exp, ca_exp = typing.cast(
+        tuple[NamedArray, NamedArray, NamedArray, NamedArray],
+        hax.unique(
+            arr2d,
+            U3,
+            axis=Height,
+            return_index=True,
+            return_inverse=True,
+            return_counts=True,
+        ),
+    )
+    assert jnp.all(va.array == va_exp.array)
+    assert jnp.all(ia.array == ia_exp.array)
+    assert jnp.all(ina.array == ina_exp.array)
+    assert jnp.all(ca.array == ca_exp.array)

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

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