haliax 1.4.dev395__tar.gz → 1.4.dev397__tar.gz

{haliax-1.4.dev395 → haliax-1.4.dev397}/.playbooks/wrap-non-named.md

@@ -47,3 +47,6 @@ Some functions need bespoke handling. For example `jnp.unique` returns several a
 
 ## Testing
 Add tests to ensure that named and unnamed calls produce the same results and that axis names are preserved or removed correctly.
+
+## Documentation
+Once your wrapper works, document it. Add `::: haliax.your_function` (replacing the name) to `docs/api.md` so users can find the new API.

{haliax-1.4.dev395 → haliax-1.4.dev397}/AGENTS.md

@@ -41,7 +41,8 @@ repository. Follow these notes when implementing new features or fixing bugs.
   `[Object 1][full.path.object1]`)
 * **Documentation**: When adding new features, ensure that the documentation is updated accordingly.
   This includes updating the Mkdocs files and any relevant docstrings. If you add a new module or
-  significant functionality, consider adding a dedicated section in the documentation.
+  significant functionality, consider adding a dedicated section in the documentation. When you
+  wrap a new JAX function, add a reference to it in `docs/api.md` so users can discover it.
 
 ## Testing
 

{haliax-1.4.dev395 → haliax-1.4.dev397}/PKG-INFO

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

@@ -256,6 +256,7 @@ These are all more or less directly from JAX's NumPy API.
 
 ### Other Operations
 
+::: haliax.bincount
 ::: haliax.clip
 ::: haliax.isclose
 ::: haliax.pad

haliax-1.4.dev397/docs/vmap.md

@@ -0,0 +1,38 @@
+## Vectorization with `haliax.vmap`
+
+`haliax.vmap` is a [`NamedArray`][haliax.NamedArray] aware wrapper around
+[`jax.vmap`][jax.vmap].  Instead of supplying positional axis numbers you pass
+the [`Axis`][haliax.Axis] (or axis name) you want to map over.  Any
+`NamedArray` containing that axis is mapped in parallel and the axis is
+reinserted in the output.  Regular JAX arrays can be mapped as well by
+providing a `default` spec or per‑argument overrides.
+
+Unlike vanilla `jax.vmap`, you may supply **one or more axes**.  When multiple
+axes are given, the function is vmapped over each axis in turn (innermost first).
+If an axis isn't already present in the array you must also specify its size,
+either by passing an `Axis` object (`Axis("batch", 4)`) or a mapping such as
+`{"batch": 4}` so the new dimension can be inserted.
+
+### Basic Example
+
+```python
+import haliax as hax
+
+Batch = hax.Axis("batch", 4)
+
+def double(x):
+    return x * 2
+
+x = hax.arange(Batch)
+y = hax.vmap(double, Batch)(x)
+```
+
+The result `y` has the same `Batch` axis as `x`, and each element was processed
+in parallel.  With JAX you would write `jax.vmap(double)(x.array)` and manually
+specify `in_axes`, but Haliax handles the axis automatically.
+
+For applying many modules in parallel see
+[`Stacked.vmap`](scan.md#apply-blocks-in-parallel-with-vmap) which builds on this
+primitive.
+
+::: haliax.vmap

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

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

{haliax-1.4.dev395 → haliax-1.4.dev397}/src/haliax/__init__.py

@@ -79,6 +79,7 @@ from .ops import (
     unique_counts,
     unique_inverse,
     unique_all,
+    bincount,
     where,
 )
 from .partitioning import auto_sharded, axis_mapping, fsdp, named_jit, shard, shard_with_axis_mapping
@@ -1053,6 +1054,7 @@ __all__ = [
     "unique_counts",
     "unique_inverse",
     "unique_all",
+    "bincount",
     "clip",
     "tril",
     "triu",

{haliax-1.4.dev395 → haliax-1.4.dev397}/src/haliax/ops.py

@@ -430,6 +430,33 @@ def unique_all(
     return values, indices, inverse, counts
 
 
+def bincount(
+    x: NamedArray,
+    Counts: Axis,
+    *,
+    weights: NamedArray | ArrayLike | None = None,
+    minlength: int = 0,
+) -> NamedArray:
+    """Named version of `jax.numpy.bincount`.
+
+    The output axis is specified by ``Counts``.
+    """
+
+    if x.ndim != 1:
+        raise ValueError("bincount only supports 1D arrays")
+
+    w_array = None
+    if weights is not None:
+        if isinstance(weights, NamedArray):
+            weights = haliax.broadcast_to(weights, x.axes)
+            w_array = weights.array
+        else:
+            w_array = jnp.asarray(weights)
+
+    result = jnp.bincount(x.array, weights=w_array, minlength=minlength, length=Counts.size)
+    return NamedArray(result, (Counts,))
+
+
 __all__ = [
     "trace",
     "where",
@@ -444,4 +471,5 @@ __all__ = [
     "unique_counts",
     "unique_inverse",
     "unique_all",
+    "bincount",
 ]

{haliax-1.4.dev395 → haliax-1.4.dev397}/tests/test_ops.py

@@ -407,3 +407,19 @@ def test_unique_shortcuts():
     assert jnp.all(ia.array == ia_exp.array)
     assert jnp.all(ina.array == ina_exp.array)
     assert jnp.all(ca.array == ca_exp.array)
+
+
+def test_bincount():
+    X = Axis("X", 6)
+    x = hax.named([0, 1, 1, 2, 3, 1], (X,))
+    B = Axis("B", 5)
+
+    out = hax.bincount(x, B)
+    expected = jnp.bincount(x.array, length=B.size)
+    assert out.axes == (B,)
+    assert jnp.all(out.array == expected)
+
+    w = hax.arange((X,), dtype=jnp.float32)
+    out_w = hax.bincount(x, B, weights=w)
+    expected_w = jnp.bincount(x.array, weights=w.array, length=B.size)
+    assert jnp.allclose(out_w.array, expected_w)

haliax-1.4.dev395/docs/vmap.md

@@ -1,9 +0,0 @@
-## Vectorization
-
-
-This primitive is also used by [`Stacked.vmap`](scan.md#apply-blocks-in-parallel-with-vmap)
-to apply an entire stack of blocks in parallel.
-
-(This is a work in progress. Please contact dlwh for more information.)
-
-::: haliax.vmap

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

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