haliax 1.4.dev397__tar.gz → 1.4.dev399__tar.gz

{haliax-1.4.dev397 → haliax-1.4.dev399}/PKG-INFO

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

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

{haliax-1.4.dev397 → haliax-1.4.dev399}/src/haliax/core.py

@@ -1361,13 +1361,26 @@ def unbind(array: NamedArray, axis: AxisSelector) -> List[NamedArray]:
     return [haliax.auto_sharded(NamedArray(a, new_axes)) for a in arrays]
 
 
-def roll(array: NamedArray, shift: Union[int, Tuple[int, ...]], axis: AxisSelection) -> NamedArray:
-    """
-    Roll an array along an axis or axes. Analogous to np.roll
+def roll(
+    array: NamedArray,
+    shift: Union[IntScalar, Tuple[int, ...], "NamedArray"],
+    axis: AxisSelection,
+) -> NamedArray:
+    """Roll an array along an axis or axes.
+
+    ``shift`` may be a scalar ``NamedArray`` in addition to an ``int`` or tuple of
+    integers.
     """
+
     axis_indices = array.axis_indices(axis)
     if axis_indices is None:
         raise ValueError(f"axis {axis} not found in {array}")
+
+    if isinstance(shift, NamedArray):
+        if shift.ndim != 0:
+            raise TypeError("shift must be a scalar NamedArray")
+        shift = shift.array
+
     return NamedArray(jnp.roll(array.array, shift, axis_indices), array.axes)
 
 

haliax-1.4.dev399/tests/test_moe_linear.py

@@ -0,0 +1,70 @@
+import jax
+import jax.random as jrandom
+from jax import numpy as jnp
+
+import haliax as hax
+from haliax.nn import MoELinear
+
+
+def _expected_moe_linear_output(moe: MoELinear, x: hax.NamedArray, group_sizes: hax.NamedArray):
+    dim_numbers = jax.lax.RaggedDotDimensionNumbers(
+        (
+            ((x.axis_indices(moe.In),), (moe.weight.axis_indices(moe.In),)),
+            ((), ()),
+        ),
+        x.axis_indices(hax.axis.without_axes(x.axes, moe.In)),
+        (moe.weight.axis_indices(moe.Experts),),
+    )
+    out_raw = jax.lax.ragged_dot_general(
+        lhs=x.array,
+        rhs=moe.weight.array,
+        group_sizes=group_sizes.array,
+        ragged_dot_dimension_numbers=dim_numbers,
+    )
+    out_axes = hax.replace_axis(x.axes, moe.In, moe.Out)
+    out = hax.named(out_raw, out_axes)
+    if moe.bias is not None:
+        out = out + moe.bias
+    return out
+
+
+def test_moe_linear_matches_ragged_dot_general():
+    B, In, Out, E = hax.make_axes(B=3, In=4, Out=5, E=2)
+    key = jrandom.PRNGKey(0)
+    moe = MoELinear.init(E, In, Out, key=key)
+
+    x = hax.random.normal(jrandom.PRNGKey(1), (B, In))
+    group_sizes = hax.named(jnp.array([2, 1], dtype=jnp.int32), (E,))
+
+    actual = moe(x, group_sizes)
+    expected = _expected_moe_linear_output(moe, x, group_sizes)
+
+    assert actual.axes == expected.axes
+    assert jnp.allclose(actual.array, expected.array, rtol=1e-5, atol=1e-5)
+
+
+def test_moe_linear_out_first_property():
+    E, In, Out = hax.make_axes(E=2, In=4, Out=3)
+    moe = MoELinear.init(E, In, Out, key=jrandom.PRNGKey(0), out_first=True)
+    assert moe.out_first
+    assert moe.weight.axes[:3] == (E, Out, In)
+
+    moe2 = MoELinear.init(E, In, Out, key=jrandom.PRNGKey(1), out_first=False)
+    assert not moe2.out_first
+    assert moe2.weight.axes[:3] == (E, In, Out)
+
+
+def test_moe_linear_gmm_matches_ragged_dot_general():
+    B, In, Out, E = hax.make_axes(B=3, In=4, Out=5, E=2)
+    moe = MoELinear.init(E, In, Out, key=jrandom.PRNGKey(0), use_gmm=True)
+
+    x = hax.random.normal(jrandom.PRNGKey(1), (B, In))
+    group_sizes = hax.named(jnp.array([2, 1], dtype=jnp.int32), (E,))
+
+    with jax.sharding.Mesh(jax.devices(), ("data",)):
+        actual = moe(x, group_sizes)
+
+    expected = _expected_moe_linear_output(moe, x, group_sizes)
+
+    assert actual.axes == expected.axes
+    assert jnp.allclose(actual.array, expected.array, rtol=1e-5, atol=1e-5)

{haliax-1.4.dev397 → haliax-1.4.dev399}/tests/test_ops.py

@@ -423,3 +423,17 @@ def test_bincount():
     out_w = hax.bincount(x, B, weights=w)
     expected_w = jnp.bincount(x.array, weights=w.array, length=B.size)
     assert jnp.allclose(out_w.array, expected_w)
+
+
+def test_roll_scalar_named_shift():
+    H = Axis("H", 4)
+    W = Axis("W", 3)
+
+    arr = hax.arange((H, W))
+    shift = hax.named(jnp.array(1), ())
+
+    rolled = hax.roll(arr, shift, H)
+    expected = jnp.roll(arr.array, shift.array, axis=0)
+
+    assert rolled.axes == arr.axes
+    assert jnp.all(rolled.array == expected)

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

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