haliax 1.4.dev398__tar.gz → 1.4.dev400__tar.gz

{haliax-1.4.dev398 → haliax-1.4.dev400}/PKG-INFO

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

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

{haliax-1.4.dev398 → haliax-1.4.dev400}/src/haliax/core.py

@@ -15,7 +15,7 @@ import numpy as np
 
 import haliax
 import haliax.axis
-from haliax.jax_utils import is_jax_array_like, is_pallas_dslice
+from haliax.jax_utils import ensure_scalar, is_jax_array_like, is_pallas_dslice
 from haliax.util import ensure_tuple
 
 from ._src.util import index_where, py_slice, slice_t
@@ -1115,9 +1115,7 @@ def updated_slice(
         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()
+            s = ensure_scalar(s, name=str(axis))
 
         array_slice_indices[axis_index] = s
         total_length = array.axes[axis_index].size
@@ -1361,13 +1359,23 @@ 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}")
+
+    shift = ensure_scalar(shift, name="shift")
+
     return NamedArray(jnp.roll(array.array, shift, axis_indices), array.axes)
 
 

{haliax-1.4.dev398 → haliax-1.4.dev400}/src/haliax/jax_utils.py

@@ -153,6 +153,21 @@ def is_scalarish(x):
         return jnp.isscalar(x) or (getattr(x, "shape", None) == ())
 
 
+def ensure_scalar(x, *, name: str = "value"):
+    """Return ``x`` if it is not a :class:`NamedArray`, otherwise ensure it is a scalar.
+
+    This is useful for APIs that can accept either Python scalars or scalar
+    ``NamedArray`` objects (for example ``roll`` or ``updated_slice``).  If ``x``
+    is a ``NamedArray`` with rank greater than 0 a :class:`TypeError` is raised.
+    """
+
+    if isinstance(x, haliax.NamedArray):
+        if x.ndim != 0:
+            raise TypeError(f"{name} must be a scalar NamedArray")
+        return x.array
+    return x
+
+
 def is_on_mac_metal():
     return jax.devices()[0].platform.lower() == "metal"
 

{haliax-1.4.dev398 → haliax-1.4.dev400}/src/haliax/ops.py

@@ -9,7 +9,7 @@ import haliax
 
 from .axis import Axis, AxisSelector, axis_name
 from .core import NamedArray, NamedOrNumeric, broadcast_arrays, broadcast_arrays_and_return_axes, named
-from .jax_utils import is_scalarish
+from .jax_utils import ensure_scalar, is_scalarish
 
 
 def trace(array: NamedArray, axis1: AxisSelector, axis2: AxisSelector, offset=0, dtype=None) -> NamedArray:
@@ -89,7 +89,7 @@ def where(
             x = named(x, ())
         x, y = broadcast_arrays(x, y)
         if isinstance(condition, NamedArray):
-            condition = condition.scalar()
+            condition = ensure_scalar(condition, name="condition")
         return jax.lax.cond(condition, lambda _: x, lambda _: y, None)
 
     condition, x, y = broadcast_arrays(condition, x, y)  # type: ignore

{haliax-1.4.dev398 → haliax-1.4.dev400}/src/haliax/wrap.py

@@ -5,7 +5,7 @@ import jax
 from haliax.core import NamedArray, _broadcast_order, broadcast_to
 
 from .axis import AxisSelection, AxisSelector, axis_spec_to_shape_dict, eliminate_axes
-from .jax_utils import is_scalarish
+from .jax_utils import ensure_scalar, is_scalarish
 
 
 def wrap_elemwise_unary(f, a, *args, **kwargs):
@@ -105,7 +105,7 @@ def wrap_elemwise_binary(op):
             else:
                 if is_scalarish(b):
                     return NamedArray(op(a.array, b), a.axes)
-                a = a.scalar()
+                a = ensure_scalar(a)
                 return op(a, b)
 
             return NamedArray(op(a.array, b), a.axes)
@@ -119,7 +119,7 @@ def wrap_elemwise_binary(op):
             else:
                 if is_scalarish(a):
                     return NamedArray(op(a, b.array), b.axes)
-                b = b.scalar()
+                b = ensure_scalar(b)
                 return op(a, b)
 
             return NamedArray(op(a, b.array), b.axes)

{haliax-1.4.dev398 → haliax-1.4.dev400}/tests/test_ops.py

@@ -423,3 +423,28 @@ 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)
+
+
+def test_roll_bad_named_shift():
+    H = Axis("H", 4)
+    W = Axis("W", 3)
+
+    arr = hax.arange((H, W))
+    shift = hax.arange((Axis("dummy", 2),))
+
+    with pytest.raises(TypeError):
+        hax.roll(arr, shift, H)

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

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