haliax 1.4.dev351__tar.gz → 1.4.dev353__tar.gz

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

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

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

{haliax-1.4.dev351 → haliax-1.4.dev353}/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)
 
 
@@ -931,6 +964,7 @@ def _compute_new_axes_and_slices_for_index(
     ordered_slices: list = [py_slice(None, None, None)] * len(array.axes)  # type: ignore
     kept_axes = [True] * len(array.axes)
     array_slice_indices = []
+    index_axis_names = set()
 
     for axis, slice_ in slices.items():
         axis_index = array._lookup_indices(axis)
@@ -946,6 +980,8 @@ def _compute_new_axes_and_slices_for_index(
             ordered_slices[axis_index] = slice_
             array_slice_indices.append(axis_index)
             kept_axes[axis_index] = False
+            for ax in slice_.axes:
+                index_axis_names.add(ax.name)
         elif isinstance(slice_, list):
             # we'll let JAX complain if this is wrong
             ordered_slices[axis_index] = slice_
@@ -955,10 +991,17 @@ def _compute_new_axes_and_slices_for_index(
                 ordered_slices[axis_index] = slice_
                 kept_axes[axis_index] = False
             elif slice_.ndim == 1:
-                # we allow this if it's a 1-d array, in which case we treat it as sugar for NamedArray(slice_, sliced-axis)
-                ordered_slices[axis_index] = haliax.named(slice_, axis_name(axis))
+                target_axis = None
+                for i2, ax2 in enumerate(array.axes):
+                    if i2 != axis_index and kept_axes[i2] and ax2.size == slice_.shape[0]:
+                        target_axis = ax2
+                        break
+                if target_axis is None:
+                    target_axis = axis
+                ordered_slices[axis_index] = haliax.named(slice_, axis_name(target_axis))
                 kept_axes[axis_index] = False
                 array_slice_indices.append(axis_index)
+                index_axis_names.add(axis_name(target_axis))
             else:
                 raise ValueError(
                     f"Only 0-d or 1-d unnamed arrays can be used for indexing. Got {slice_} for axis {axis}"
@@ -966,6 +1009,14 @@ def _compute_new_axes_and_slices_for_index(
         else:
             raise ValueError(f"Only NamedArrays can be used for advanced indexing. Got {slice_} for axis {axis}")
 
+    # If any index array uses axes that are already present in the array and not removed,
+    # we need to explicitly advance-index those axes so numpy broadcasting works.
+    for i, ax in enumerate(array.axes):
+        if kept_axes[i] and ax.name in index_axis_names:
+            ordered_slices[i] = haliax.arange(ax)
+            array_slice_indices.append(i)
+            kept_axes[i] = False
+
     # advanced indexing
     if len(array_slice_indices) > 0:
         # this requires broadcasting

{haliax-1.4.dev351 → haliax-1.4.dev353}/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.dev353/tests/test_scatter_gather.py

@@ -0,0 +1,165 @@
+import jax
+import jax.numpy as jnp
+
+import haliax as hax
+from haliax import Axis, dslice
+
+
+# -----------------------------------------------------------------------------
+# Helper for reference via take_along_axis
+# -----------------------------------------------------------------------------
+
+
+def _ref_gather(src, axis, idx):
+    ax_num = src.axes.index(axis)
+    # broadcast idx to match src without the gathered axis
+    other_axes = tuple(ax for ax in src.axes if ax != axis)
+    broadcast_axes = other_axes
+    for ax in idx.axes:
+        if ax not in broadcast_axes:
+            broadcast_axes += (ax,)
+    idx_b = hax.broadcast_to(idx, broadcast_axes, enforce_no_extra_axes=False)
+    if idx_b.array.ndim == src.array.ndim - 1:
+        idx_arr = idx_b.array[..., None]
+    else:
+        idx_arr = idx_b.array
+    out = jnp.take_along_axis(src.array, idx_arr, axis=ax_num)
+    if idx_b.array.ndim == src.array.ndim - 1:
+        out = out.squeeze(ax_num)
+    return out
+
+
+# ---------------------------- 1. single batched selector ----------------------
+
+
+def test_single_batched_selector():
+    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]
+    assert out.axes == (B, S)
+    assert jnp.array_equal(out.array, _ref_gather(x, V, idx))
+
+
+# ---------------------------- 2. selector adds new axis -----------------------
+
+
+def test_selector_adds_new_axis():
+    B, S, V, T = Axis("batch", 2), Axis("seq", 3), Axis("vocab", 5), Axis("step", 4)
+    logits = hax.arange((B, S, V))
+    idx = hax.arange((B, T), dtype=jnp.int32) % V.size
+    out = logits["vocab", idx]
+    assert set(out.axes) == {B, S, T}
+    ref = jnp.transpose(_ref_gather(logits, V, idx), (0, 2, 1))
+    assert jnp.array_equal(out.array, ref)
+
+
+# ------------------------ 3. two contiguous selector arrays -------------------
+
+
+def test_two_contiguous_selectors():
+    B, X, Y = Axis("batch", 3), Axis("x", 5), Axis("y", 7)
+    a = hax.arange((B, X, Y))
+    ix = hax.arange((B,), dtype=jnp.int32) % X.size
+    iy = hax.arange((B,), dtype=jnp.int32) % Y.size
+    out = a["x", ix, "y", iy]
+    assert out.axes == (B,)
+    ref = a.array[jnp.arange(3), ix.array, iy.array]
+    assert jnp.array_equal(out.array, ref)
+
+
+# ------------------ 4. non-contiguous selectors → axes to front --------------
+
+
+def test_noncontig_selectors():
+    B, X, Z, Y = Axis("batch", 2), Axis("x", 4), Axis("z", 6), Axis("y", 5)
+    a = hax.arange((B, X, Z, Y))
+    ix = hax.arange((B,), dtype=jnp.int32) % X.size
+    iy = hax.arange((B,), dtype=jnp.int32) % Y.size
+    out = a["x", ix, "y", iy]
+    assert out.axes == (B, Z)
+    ref = a.array[jnp.arange(2), ix.array, :, iy.array]
+    assert jnp.array_equal(out.array, ref)
+
+
+# ----------------- 5. integer elimination + selector --------------------------
+
+
+def test_mixed_int_and_selector():
+    B, C, V = Axis("batch", 3), Axis("channel", 2), Axis("vocab", 6)
+    x = hax.arange((B, C, V))
+    idx = hax.arange((B,), dtype=jnp.int32) % V.size
+    out = x["channel", 1, "vocab", idx]
+    assert out.axes == (B,)
+    ref = x.array[:, 1, :][jnp.arange(3), idx.array]
+    assert jnp.array_equal(out.array, ref)
+
+
+def test_dslice_with_selector():
+    B, S, V = Axis("batch", 2), Axis("seq", 5), Axis("vocab", 10)
+    x = hax.arange((B, S, V))
+    idx = (hax.arange((B, S), dtype=jnp.int32) + 2) % 4
+    shard = V.resize(4)
+    x_shard = x["vocab", dslice(0, shard)]
+    out = x_shard["vocab", idx]
+    assert out.axes == (B, S)
+    ref = x.array[:, :, :4][jnp.arange(B.size)[:, None], jnp.arange(S.size)[None, :], idx.array]
+    assert jnp.array_equal(out.array, ref)
+
+
+def test_scalar_eliminates_axis():
+    B, S, V = Axis("batch", 2), Axis("seq", 3), Axis("vocab", 4)
+    x = hax.arange((B, S, V))
+    out = x["seq", 1]
+    assert out.axes == (B, V)
+    assert jnp.array_equal(out.array, x.array[:, 1, :])
+
+
+# ----------------- 9. plain ndarray selector sugar ----------------------------
+
+
+def test_plain_ndarray_selector():
+    B, V = Axis("batch", 3), Axis("vocab", 5)
+    x = hax.arange((B, V))
+    idx = jnp.array([0, 2, 4], dtype=jnp.int32)
+    out = x["vocab", idx]
+    assert out.axes == (B,)
+    assert jnp.array_equal(out.array, x.array[jnp.arange(3), idx])
+
+
+# ----------------- 10. two selectors needing broadcast ------------------------
+
+
+def test_multiselector_broadcast():
+    B, S, V = Axis("batch", 2), Axis("seq", 3), Axis("vocab", 6)
+    a = hax.arange((B, S, V))
+    idx1 = hax.arange((B, S), dtype=jnp.int32) % V.size
+    out = a["vocab", idx1]
+    assert out.axes == (B, S)
+    assert jnp.array_equal(out.array, _ref_gather(a, V, idx1))
+
+
+# ----------------- 11. scatter-ADD via .at[…].add -----------------------------
+
+
+def test_scatter_add():
+    B, S, V = Axis("batch", 2), Axis("seq", 3), Axis("vocab", 5)
+    x = hax.zeros((B, S, V))
+    idx = hax.arange((B, S), dtype=jnp.int32) % V.size
+    ones = hax.ones((B, S))
+    y = x.at[{V: idx}].add(ones)
+    ref = jnp.zeros((2, 3, 5)).at[jnp.arange(2)[:, None], jnp.arange(3)[None, :], idx.array].add(1.0)
+    assert jnp.array_equal(y.array, ref)
+
+
+# ----------------- 12. scatter-SET via .at[…].set -----------------------------
+
+
+def test_scatter_set():
+    B, V = Axis("batch", 2), Axis("vocab", 6)
+    x = hax.zeros((B, V))
+    idx = hax.named(jnp.array([1, 4]), B)
+    val = hax.ones(B) * 9
+    y = x.at[{V: idx}].set(val)
+    ref = jnp.zeros((2, 6)).at[jnp.arange(2), idx.array].set(9)
+    assert jnp.array_equal(y.array, ref)

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

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