haliax 1.4.dev353__tar.gz → 1.4.dev354__tar.gz

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

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

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

{haliax-1.4.dev353 → haliax-1.4.dev354}/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.dev354}/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"