haliax 1.4.dev353__tar.gz → 1.4.dev355__tar.gz

{haliax-1.4.dev353 → haliax-1.4.dev355}/PKG-INFO

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

@@ -287,3 +287,42 @@ operation more effectively.)
 
     It's worth emphasizing that these functions are typically compiled to scatter-add and friends (as appropriate).
     This is the preferred way to do scatter/gather operations in JAX, as well as in Haliax.
+
+## Scatter/Gather
+
+Haliax supports scatter/gather semantics in its indexing operations. When an axis
+is indexed by another NamedArray (or a 1-D JAX array), the values of that axis
+are gathered according to the index array and the axes of the indexer are
+inserted into the result.
+
+```python
+import haliax as hax
+import jax.numpy as jnp
+
+B, S, V = Axis("batch", 4), Axis("seq", 3), Axis("vocab", 7)
+x = hax.arange((B, S, V))
+idx = hax.arange((B, S), dtype=jnp.int32) % V.size
+
+out = x["vocab", idx]
+```
+
+Here `out` has axes `(B, S)` and its values match `jax.numpy.take_along_axis`
+on the underlying ndarray.
+
+For scatter-style updates where each batch writes to a different position, use
+[`updated_slice`][haliax.updated_slice]:
+
+```python
+Batch = hax.Axis("batch", 2)
+Seq = hax.Axis("seq", 5)
+New = hax.Axis("seq", 2)
+
+cache = hax.zeros((Batch, Seq), dtype=int)
+lengths = hax.named([1, 3], axis=Batch)
+kv = hax.named([[1, 2], [3, 4]], axis=(Batch, New))
+
+result = updated_slice(cache, {"seq": lengths}, kv)
+```
+
+This inserts `[1, 2]` starting at position `1` in batch `0` and `[3, 4]` starting
+at position `3` in batch `1`.

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

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

{haliax-1.4.dev353 → haliax-1.4.dev355}/src/haliax/hof.py

@@ -19,16 +19,16 @@ from ._src.scan import (
     map,
     scan,
 )
-from .axis import Axis, AxisSelector, selects_axis
+from .axis import Axis, AxisSelection, AxisSelector, selects_axis
 from .core import NamedArray
 from .jax_utils import Static, broadcast_prefix, is_jax_array_like
 from .partitioning import physical_axis_name
-from .util import is_named_array
+from .util import ensure_tuple, is_named_array
 
 
 def vmap(
     fn,
-    axis: AxisSelector,
+    axis: AxisSelection,
     *,
     default: PyTree[UnnamedAxisSpec] = _zero_if_array_else_none,
     args: PyTree[UnnamedAxisSpec] = (),
@@ -43,7 +43,9 @@ def vmap(
 
     Args:
         fn (Callable): function to vmap over
-        axis (Axis): axis to vmap over
+        axis (Axis or Sequence[Axis]): axis or axes to vmap over. If a sequence is
+            provided, the function will be vmapped over each axis in turn,
+            from innermost to outermost.
         default: how to handle (unnamed) arrays by default. Should be either an integer or None, or a callable that takes a PyTree leaf
             and returns an integer or None, or a PyTree prefix of the same. If an integer, the array will be mapped over that axis. If None, the array will not be mapped over.
         args: optional per-argument overrides for how to handle arrays. Should be a PyTree prefix of the same type as default.
@@ -53,6 +55,15 @@ def vmap(
     if kwargs is None:
         kwargs = {}
 
+    axes = ensure_tuple(axis)  # type: ignore
+    if len(axes) > 1:
+        mapped = fn
+        for ax in reversed(axes):
+            mapped = vmap(mapped, ax, default=default, args=args, kwargs=kwargs)
+        return mapped
+    else:
+        axis = axes[0]
+
     signature = inspect.signature(fn)
 
     # this mirrors equinox's filter_vmap, but it's not really documented there so:

{haliax-1.4.dev353 → haliax-1.4.dev355}/tests/test_hof.py

@@ -321,3 +321,22 @@ def test_vmap_error_for_incorrectly_specified_args():
     Width = Axis("Width", 3)
 
     hax.vmap(lambda a: Module(a), Batch)(Width)
+
+
+def test_vmap_multiple_axes():
+    Batch1 = Axis("Batch1", 4)
+    Batch2 = Axis("Batch2", 3)
+    Width = Axis("Width", 2)
+    Depth = Axis("Depth", 5)
+
+    named = hax.random.uniform(PRNGKey(0), (Batch1, Batch2, Width, Depth))
+
+    def vmap_fun(x):
+        return x.sum(Width)
+
+    selected = hax.vmap(vmap_fun, (Batch1, Batch2))(named)
+
+    expected = jnp.sum(named.array, axis=2)
+
+    assert jnp.allclose(selected.array, expected)
+    assert selected.axes == (Batch1, Batch2, Depth)

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

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