haliax 1.4.dev352__tar.gz → 1.4.dev354__tar.gz

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

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

@@ -1,7 +1,6 @@
 import contextlib
 import functools as ft
 import typing
-import warnings
 from dataclasses import dataclass
 from math import prod
 from types import EllipsisType
@@ -324,7 +323,9 @@ class NamedArray:
     def slice(self, *args, **kwargs) -> "NamedArray":
         return haliax.slice(self, *args, **kwargs)
 
-    def updated_slice(self, start: Mapping[AxisSelector, int], update: "NamedArray") -> "NamedArray":
+    def updated_slice(
+        self, start: Mapping[AxisSelector, Union[int, "NamedArray"]], update: "NamedArray"
+    ) -> "NamedArray":
         return haliax.updated_slice(self, start=start, update=update)
 
     def take(self, axis: AxisSelector, index: Union[int, "NamedArray"]) -> "NamedArray":
@@ -854,7 +855,7 @@ def _slice_new(
 
 
 def updated_slice(
-    array: NamedArray, start: Mapping[AxisSelector, Union[int, jnp.ndarray]], update: NamedArray
+    array: NamedArray, start: Mapping[AxisSelector, Union[int, jnp.ndarray, NamedArray]], update: NamedArray
 ) -> NamedArray:
     """
     Updates a slice of an array with another array.
@@ -868,18 +869,42 @@ def updated_slice(
         NamedArray: The updated array.
     """
 
+    # figure out which axis‐names to map over
+    map_axes: list[str] = []
+    for axis_sel, s in start.items():
+        if isinstance(s, NamedArray):
+            for ax in s.axes:
+                if ax.name not in map_axes:
+                    map_axes.append(ax.name)
+
+    # need to vmap
+    if len(map_axes) > 0:
+        # scalar version: all starts are ints / tracers
+        f = updated_slice
+        for axis_name in map_axes:
+            # make sure that axis_name is in `array`. otherwise it doesn't make sense to vmap over it
+            if array._lookup_indices(axis_name) is None:
+                raise ValueError(f"axis {axis_name} not found in original array's axes: {array.shape}")
+            f = haliax.vmap(f, axis=axis_name)
+        return f(array, start, update)
+
     array_slice_indices = [0] * len(array.axes)
     for axis, s in start.items():
-        axis_index = array._lookup_indices(axis_name(axis))
+        axis_index = array._lookup_indices(haliax.axis_name(axis))
         if axis_index is None:
             raise ValueError(f"axis {axis} not found in {array}")
+        if isinstance(s, NamedArray):  # this can happen in the vmap case
+            if s.ndim != 0:
+                raise ValueError(f"NamedArray {s} must be a scalar for axis {axis} in updated_slice")
+            s = s.scalar()
+
         array_slice_indices[axis_index] = s
         total_length = array.axes[axis_index].size
-        update_axis = update._lookup_indices(axis_name(axis))
+        update_axis = update._lookup_indices(haliax.axis_name(axis))
 
-        if update_axis is None:
-            raise ValueError(f"axis {axis} not found in {update}")
         # if s is a tracer we can't check the size
+        if update_axis is None:
+            continue
         if isinstance(s, int) and update.axes[update_axis].size + s > total_length:
             raise ValueError(
                 f"update axis {axis} is too large to start at {s}. Array size is {total_length}, update size is"
@@ -888,17 +913,25 @@ def updated_slice(
 
     # broadcasting here is a bit delicate because the sizes aren't necessarily the same
     # we need to broadcast the update array to the same axis names as the array we're updating, adding them as necessary
-    broadcasted_axes = []
-    for axis in array.axes:
-        update_axis = update._lookup_indices(axis.name)
-        if update_axis is None:
-            broadcasted_axes.append(axis)
-        else:
-            broadcasted_axes.append(update.axes[update_axis])
-
-    update = haliax.broadcast_to(update, broadcasted_axes, enforce_no_extra_axes=True)
+    if update.ndim > 0:
+        # if there are any axes in update that are not in array, it is an error:
+        axes_in_update = haliax.axis.without_axes(update.axes, array.axes)
+        if axes_in_update:
+            raise ValueError(
+                f"Update array with shape {update.shape} has axes {axes_in_update} that are not in the original array"
+                f" with shape {array.shape}. This is not allowed in updated_slice."
+            )
+        broadcasted_axes = []
+        for ax in array.axes:
+            upd_ax = update._lookup_indices(ax.name)
+            broadcasted_axes.append(ax if upd_ax is None else update.axes[upd_ax])
+        update = haliax.broadcast_to(update, broadcasted_axes, enforce_no_extra_axes=True)
+        upd_arr = update.array
+    else:
+        # scalar case: just add one axis so it doesn't get too mad
+        upd_arr = update.array.reshape((1,))
 
-    updated = jax.lax.dynamic_update_slice(array.array, update.array, array_slice_indices)
+    updated = jax.lax.dynamic_update_slice(array.array, upd_arr, array_slice_indices)
     return NamedArray(updated, array.axes)
 
 

{haliax-1.4.dev352 → 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.dev352 → haliax-1.4.dev354}/tests/core_test.py

@@ -1,11 +1,8 @@
-import equinox as eqx
 import jax
 import jax.numpy as jnp
-import pytest
 from jax.random import PRNGKey
 
-import haliax as hax
-from haliax import Axis, NamedArray
+from haliax import Axis, NamedArray, updated_slice
 
 
 def test_unary_np_functions():
@@ -603,3 +600,90 @@ def test_slice_nd_array_present_dims():
     assert jnp.all(jnp.equal(named1[{"H": index1}].array, named1.array[index1.array, :, :]))
 
     # this is not ok, since the H would not be eliminated
+
+
+import numpy as np
+import pytest
+
+import haliax as hax
+
+
+def test_scalar_updated_slice():
+    # Base case: scalar start on a 1D array
+    Seq = hax.Axis("seq", 5)
+    arr = hax.arange((Seq,), dtype=int)
+    # replace positions 2 and 3 with [100, 101]
+    upd = hax.named([100, 101], "seq")
+
+    result = updated_slice(arr, {"seq": 2}, upd)
+    # expect [0,1,100,101,4]
+    assert np.array_equal(result.array, np.array([0, 1, 100, 101, 4]))
+
+
+def test_ragged_single_token():
+    # Ragged case: one token per batch at different positions
+    Batch = hax.Axis("batch", 3)
+    Seq = hax.Axis("seq", 5)
+    cache = hax.zeros((Batch, Seq), dtype=int)
+
+    # lengths[b] is next free slot for batch b
+    lengths = hax.named([0, 1, 2], axis=Batch)
+    new_k = hax.named([7, 8, 9], axis=Batch)
+
+    result = updated_slice(cache, {"seq": lengths}, new_k)
+
+    # build expected NumPy array
+    exp = np.zeros((3, 5), int)
+    exp[0, 0] = 7
+    exp[1, 1] = 8
+    exp[2, 2] = 9
+
+    assert np.array_equal(result.array, exp)
+
+
+def test_ragged_multi_token():
+    # Ragged case: a block of 2 tokens per batch at different positions
+    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)
+    # for batch=0 insert [1,2] at pos=1, for batch=1 insert [3,4] at pos=3
+    kv = hax.named([[1, 2], [3, 4]], axis=(Batch, New))
+
+    result = updated_slice(cache, {"seq": lengths}, kv)
+
+    exp = np.zeros((2, 5), int)
+    exp[0, 1] = 1
+    exp[0, 2] = 2
+    exp[1, 3] = 3
+    exp[1, 4] = 4
+
+    assert np.array_equal(result.array, exp)
+
+
+def test_ragged_multi_token_bad_axis_name():
+    # Ragged case: a block of 2 tokens per batch at different positions
+    Batch = hax.Axis("batch", 2)
+    Seq = hax.Axis("seq", 5)
+    New = hax.Axis("new", 2)
+
+    cache = hax.zeros((Batch, Seq), dtype=int)
+    lengths = hax.named([1, 3], axis=Batch)
+    # for batch=0 insert [1,2] at pos=1, for batch=1 insert [3,4] at pos=3
+    kv = hax.named([[1, 2], [3, 4]], axis=(Batch, New))
+
+    with pytest.raises(ValueError, match="that are not in the original array with shape "):
+        updated_slice(cache, {"seq": lengths}, kv)
+
+
+def test_update_overflow_error():
+    # Overflow: scalar start + update too large for axis → ValueError
+    Seq = hax.Axis("seq", 4)
+    arr = hax.zeros((Seq,), dtype=int)
+    # update of length 3 starting at pos=2 would run off the end (2+3 > 4)
+    upd = hax.arange((hax.Axis("seq", 3),), dtype=int)
+
+    with pytest.raises(ValueError):
+        updated_slice(arr, {"seq": 2}, upd)

{haliax-1.4.dev352 → 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.dev352/src/haliax/__about__.py

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