haliax 1.4.dev386__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.dev386 → 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.dev386 → haliax-1.4.dev388}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: haliax
-Version: 1.4.dev386
+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.dev386 → 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, unique, 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
@@ -1035,6 +1049,10 @@ __all__ = [
     "trace",
     "where",
     "unique",
+    "unique_values",
+    "unique_counts",
+    "unique_inverse",
+    "unique_all",
     "clip",
     "tril",
     "triu",

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

@@ -342,4 +342,106 @@ def unique(
     return tuple(ret)
 
 
-__all__ = ["trace", "where", "tril", "triu", "isclose", "pad_left", "pad", "clip", "unique"]
+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.dev386 → 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
@@ -359,3 +360,50 @@ def test_unique():
 
     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.dev386/src/haliax/__about__.py

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