haliax 1.4.dev348__tar.gz → 1.4.dev352__tar.gz

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

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

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

{haliax-1.4.dev348 → haliax-1.4.dev352}/src/haliax/_src/state_dict.py

@@ -215,6 +215,8 @@ def from_state_dict(tree: T, state_dict: StateDict, prefix: Optional[str] = None
             raise ValueError("Cannot extract a leaf value from a state dict without a prefix")
         # TODO: add "strict" flag so we can return None in cases where it's just missing
         return jnp.array(state_dict[prefix])
+    elif tree is None:
+        return None
     else:
         if prefix is None:
             return tree

{haliax-1.4.dev348 → haliax-1.4.dev352}/src/haliax/core.py

@@ -931,6 +931,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 +947,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 +958,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 +976,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.dev352/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.dev348 → haliax-1.4.dev352}/tests/test_state_dict.py

@@ -175,3 +175,27 @@ def test_stacked_layer_norm():
     new_norms = unflatten_modules_from_export(norms_flat, norms2)
 
     assert norms == new_norms
+
+
+def test_linear_doesnt_read_bias_if_it_didnt_have_bias():
+    H = hax.Axis("H", 10)
+    W = hax.Axis("W", 20)
+    D = hax.Axis("D", 30)
+    B = hax.Axis("B", 40)
+
+    linear = hax.nn.Linear.init((H, W), (D, B), key=jax.random.PRNGKey(0), use_bias=False, out_first=True)
+
+    flat_linear = linear.flatten_for_export()
+
+    flat_state_dict = to_state_dict(flat_linear)
+
+    assert "bias" not in flat_state_dict
+    flat_state_dict["bias"] = jnp.zeros((D.size * B.size,))  # add a dummy bias
+
+    # now unflatten it
+    linear2 = Linear.init((H, W), (D, B), key=jax.random.PRNGKey(1), use_bias=False, out_first=True)
+    flinear2 = linear2.flatten_for_export()
+    flinear2 = from_state_dict(flinear2, flat_state_dict)
+    new_linear = flinear2.unflatten_from_export(linear2)
+
+    assert linear == new_linear

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

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