haliax 1.4.dev379__tar.gz → 1.4.dev381__tar.gz

{haliax-1.4.dev379 → haliax-1.4.dev381}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: haliax
-Version: 1.4.dev379
+Version: 1.4.dev381
 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.dev379 → haliax-1.4.dev381}/docs/api.md

@@ -258,6 +258,7 @@ These are all more or less directly from JAX's NumPy API.
 
 ::: haliax.clip
 ::: haliax.isclose
+::: haliax.pad
 ::: haliax.top_k
 ::: haliax.trace
 ::: haliax.tril

{haliax-1.4.dev379 → haliax-1.4.dev381}/docs/indexing.md

@@ -147,6 +147,11 @@ def f(x, slice_size: int):
 f(q, 2)
 ```
 
+When indexing with ``dslice`` the slice is gathered starting at ``start`` for
+``size`` elements.  Reads beyond the end of the array return the ``fill_value``
+(0 by default).  When used with ``at`` updates, any writes outside the bounds of
+the array are dropped.  These semantics match JAX's scatter/gather behavior.
+
 For convenience/brevity, `dslice` is aliased as `ds`. In addition, we also expose `dblock`, which is a convenience
 function for computing the start and size of a slice given a block index and the size of the slice. Thus, the above
 example can be written as follows:

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

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

{haliax-1.4.dev379 → haliax-1.4.dev381}/src/haliax/__init__.py

@@ -66,7 +66,7 @@ from .core import (
 from .haxtyping import Named
 from .hof import fold, map, scan, vmap
 from .jax_utils import tree_checkpoint_name
-from .ops import clip, isclose, pad_left, trace, tril, triu, where
+from .ops import clip, isclose, pad_left, pad, trace, tril, triu, where
 from .partitioning import auto_sharded, axis_mapping, fsdp, named_jit, shard, shard_with_axis_mapping
 from .specialized_fns import top_k
 from .types import Scalar
@@ -1082,6 +1082,7 @@ __all__ = [
     "are_shape_checks_enabled",
     "isclose",
     "pad_left",
+    "pad",
     "stack",
     "concatenate",
     "eliminate_axes",

{haliax-1.4.dev379 → haliax-1.4.dev381}/src/haliax/axis.py

@@ -598,14 +598,19 @@ def resolve_axis(axis_spec: AxisSpec, axis_selection: AxisSelection) -> AxisSpec
 
 
 class dslice(eqx.Module):
-    """
-    Dynamic slice, comprising a (start, length) pair. Also aliased as ds.
+    """Dynamic slice, comprising a (start, length) pair. Also aliased as ``ds``.
+
+    NumPy-style slices like ``a[i:i+16]`` don't work inside :func:`jax.jit`, because
+    JAX requires slice bounds to be static. ``dslice`` works around this by
+    separating the dynamic ``start`` from the static ``size`` so that you can
+    write ``a[dslice(i, 16)]`` or simply ``a[ds(i, 16)]``.
 
-    Normal numpy-isms like a[i:i+16] don't work in Jax inside jit, because slice doesn't like tracers and JAX
-    can't see that the slice is constant. This is a workaround that lets you do a[dslice(i, 16)] or even a[ds(i, 16)]
-    instead.
+    When used in indexing or ``at`` updates, ``dslice`` behaves like a gather of
+    ``size`` elements starting at ``start``. Reads beyond the end of the array are
+    filled with a value (0 by default) and writes outside the array bounds are
+    dropped, matching JAX's default scatter/gather semantics.
 
-    This class's name is taken from [jax.experimental.pallas.dslice][].
+    This class's name is taken from :mod:`jax.experimental.pallas`.
     """
 
     start: int

{haliax-1.4.dev379 → haliax-1.4.dev381}/src/haliax/core.py

@@ -1163,6 +1163,11 @@ def index(array: NamedArray, slices: Mapping[AxisSelector, NamedIndex]) -> Named
     you might use `array[{"batch": slice(0, 10)}]` or `array["batch", 0:10]` to select the first 10 elements
     of the 'batch' axis.
 
+    When indexing with a ``dslice`` the slice is gathered starting at the given
+    ``start`` for ``size`` elements. Values read past the end of the array are
+    filled with the ``fill_value`` (defaults to ``0``), and writes outside the
+    bounds are dropped.
+
     See Also:
         * [haliax.NamedArray.at][] for a functional equivalent of in-place array modifications.
 
@@ -1172,8 +1177,12 @@ def index(array: NamedArray, slices: Mapping[AxisSelector, NamedIndex]) -> Named
     """
     # indices where we have array args
     new_axes, ordered_slices = _compute_new_axes_and_slices_for_index(array, slices)
-    sliced, ordered_slices = _handle_dynamic_slices(array.array, ordered_slices)
-    sliced = sliced[tuple(ordered_slices)]
+    needs_fill = any(isinstance(s, dslice) or is_pallas_dslice(s) for s in slices.values())
+    ordered_slices = [s.array if isinstance(s, NamedArray) else s for s in ordered_slices]
+    if needs_fill:
+        sliced = array.array.at[tuple(ordered_slices)].get(mode="fill", fill_value=0)
+    else:
+        sliced = array.array[tuple(ordered_slices)]
 
     return haliax.named(sliced, new_axes)
 
@@ -1190,9 +1199,19 @@ def _compute_new_axes_and_slices_for_index(
         axis_index = array.axis_indices(axis)
         if axis_index is None:
             raise ValueError(f"axis {axis} not found in {array}")
-        if isinstance(slice_, py_slice) or isinstance(slice_, dslice) or is_pallas_dslice(slice_):
+        if isinstance(slice_, py_slice):
             ordered_slices[axis_index] = slice_
             kept_axes[axis_index] = True
+        elif is_pallas_dslice(slice_) or isinstance(slice_, dslice):
+            # we want to treat this like a slice, but the closest approximation is an arange, but those have to broadcast.
+            # So instead we make a named array arange with the same name as the axis.
+            start = slice_.start
+            size = slice_.size
+            orig_axis = array.axes[axis_index]
+            ordered_slices[axis_index] = haliax.arange(orig_axis.resize(size), start=start)
+            kept_axes[axis_index] = False
+            array_slice_indices.append(axis_index)
+            index_axis_names.add(orig_axis.name)
         elif isinstance(slice_, int):
             ordered_slices[axis_index] = slice_
             kept_axes[axis_index] = False
@@ -1292,34 +1311,6 @@ def _compute_new_axes_and_slices_for_index(
     return new_axes, ordered_slices
 
 
-def _handle_dynamic_slices(array: jnp.ndarray, slices):
-    """
-    Helper function to handle dynamic slices in the array. These have to be handled with jax.lax.dynamic_slice,
-    which is for when the start index is not known at compile time. (Sizes must always be known at compile time.)
-
-    Notes:
-        **MUTATES `slices` IN PLACE**
-
-    Returns:
-        array.array: the sliced array
-
-    """
-    indices_for_dslice = [0] * array.ndim
-    lengths_for_dslice = list(array.shape)
-    dslice_indices = []
-    need_to_slice = False
-    for axis_index, slice_ in enumerate(slices):
-        if isinstance(slice_, dslice) or is_pallas_dslice(slice_):
-            dslice_indices.append(axis_index)
-            indices_for_dslice[axis_index] = slice_.start
-            lengths_for_dslice[axis_index] = slice_.size
-            need_to_slice = True
-    if need_to_slice:
-        array = jax.lax.dynamic_slice(array, indices_for_dslice, lengths_for_dslice)
-        for i in dslice_indices:
-            slices[i] = py_slice(None, None, None)
-    return array, slices
-
 
 def split(a: NamedArray, axis: AxisSelector, new_axes: Sequence[Axis]) -> Sequence[NamedArray]:
     """
@@ -2056,38 +2047,8 @@ class _NamedIndexUpdateRef:
 
 def _raw_indices_for_at(array, indexes):
     sliced_axes, ordered_slices = _compute_new_axes_and_slices_for_index(array, indexes)
-    del sliced_axes
-    # this isn't the fastest (it does the _compute_new_axes_and_slices_for_index twice)
-    # but it's easy
     _sliced = index(array, indexes)
-    # we have to handle dslices differently than for normal indexing, because we can't use
-    # extra dynamic_slices...
-    # we'd like to just replace these with iota, but we have account for broadcasting semantics
-    # for the other arrays
-    dslice_sizes = tuple(x.size for x in ordered_slices if isinstance(x, dslice) or is_pallas_dslice(x))  # type: ignore
-    current_array_slice_shape = next((x.shape for x in ordered_slices if is_jax_array_like(x)), None)  # type: ignore
-    dims_to_expand = list(range(len(dslice_sizes)))
-    if current_array_slice_shape is not None:
-        iota_shape = dslice_sizes + current_array_slice_shape
-    else:
-        iota_shape = dslice_sizes
-
-    def iota_for_dslice(dslice, cur_dynamic_slice):
-        return jax.lax.broadcasted_iota(int, iota_shape, cur_dynamic_slice) + dslice.start
-
-    if len(dslice_sizes) > 0:
-        cur_dynamic_slice = 0
-        for i in range(len(ordered_slices)):
-            if isinstance(ordered_slices[i], dslice) or is_pallas_dslice(ordered_slices[i]):
-                ordered_slices[i] = iota_for_dslice(ordered_slices[i], cur_dynamic_slice)
-                cur_dynamic_slice += 1
-            elif is_jax_array_like(ordered_slices[i]):
-                # prepend array slices with one 1 for each dynamic slice
-                ordered_slices[i] = jnp.expand_dims(ordered_slices[i], axis=dims_to_expand)
-
-        assert cur_dynamic_slice == len(dslice_sizes)
-
-    # ok the ordered slices are now correct
+    ordered_slices = [s.array if isinstance(s, NamedArray) else s for s in ordered_slices]
     return ordered_slices, _sliced.axes
 
 

{haliax-1.4.dev379 → haliax-1.4.dev381}/src/haliax/ops.py

@@ -1,10 +1,10 @@
 import typing
-from typing import Optional, Union
+from typing import Mapping, Optional, Union
 
 import jax
 import jax.numpy as jnp
 
-from .axis import Axis, AxisSelector
+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
 
@@ -152,10 +152,48 @@ def pad_left(array: NamedArray, axis: Axis, new_axis: Axis, value=0) -> NamedArr
     return NamedArray(padded, array.axes[:idx] + (new_axis,) + array.axes[idx + 1 :])
 
 
+def pad(
+    array: NamedArray,
+    pad_width: Mapping[AxisSelector, tuple[int, int]],
+    *,
+    mode: str = "constant",
+    constant_values: NamedOrNumeric = 0,
+    **kwargs,
+) -> NamedArray:
+    """Version of ``jax.numpy.pad`` that works with ``NamedArray``.
+
+    ``pad_width`` should be a mapping from axis (or axis name) to a ``(before, after)``
+    tuple specifying how much padding to add on each side of that axis. Any axis
+    not present in ``pad_width`` will not be padded.
+    """
+
+    padding = []
+    new_axes = []
+    for ax in array.axes:
+        left_right = pad_width.get(ax)
+        if left_right is None:
+            left_right = pad_width.get(axis_name(ax))  # type: ignore[arg-type]
+        if left_right is None:
+            left_right = (0, 0)
+        left, right = left_right
+        padding.append((left, right))
+        new_axes.append(ax.resize(ax.size + left + right))
+
+    result = jnp.pad(
+        array.array,
+        padding,
+        mode=mode,
+        constant_values=raw_array_or_scalar(constant_values),
+        **kwargs,
+    )
+
+    return NamedArray(result, tuple(new_axes))
+
+
 def raw_array_or_scalar(x: NamedOrNumeric):
     if isinstance(x, NamedArray):
         return x.array
     return x
 
 
-__all__ = ["trace", "where", "tril", "triu", "isclose", "pad_left", "clip"]
+__all__ = ["trace", "where", "tril", "triu", "isclose", "pad_left", "pad", "clip"]

{haliax-1.4.dev379 → haliax-1.4.dev381}/tests/core_test.py

@@ -586,11 +586,36 @@ def test_slice_nd_dslice():
     named1 = hax.random.uniform(PRNGKey(0), (H, W, D))
     from haliax import ds
 
-    assert jnp.all(jnp.equal(named1["H", ds(0, 5), "D", ds(3, 7)].array, named1.array[0:5, :, 3:10]))
+    ref = jnp.take(named1.array, jnp.arange(0, 5), axis=0, mode="fill", fill_value=0)
+    ref = jnp.take(ref, jnp.arange(3, 10), axis=2, mode="fill", fill_value=0)
+    ref = jnp.transpose(ref, (0, 2, 1))
+    assert jnp.all(jnp.equal(named1["H", ds(0, 5), "D", ds(3, 7)].array, ref))
     # test mixed normal and dslice
-    assert jnp.all(jnp.equal(named1["H", ds(1, 5), "D", 3:7].array, named1.array[1:6, :, 3:7]))
-    assert jnp.all(jnp.equal(named1["H", ds(2, 5), "D", 3].array, named1.array[2:7, :, 3]))
-    assert jnp.all(jnp.equal(named1["H", ds(3, 5), "D", 3:10:2].array, named1.array[3:8, :, 3:10:2]))
+    ref = jnp.take(named1.array, jnp.arange(1, 6), axis=0, mode="fill", fill_value=0)
+    ref = jnp.take(ref, jnp.arange(3, 7), axis=2, mode="fill", fill_value=0)
+    assert jnp.all(jnp.equal(named1["H", ds(1, 5), "D", 3:7].array, ref))
+    ref = jnp.take(named1.array, jnp.arange(2, 7), axis=0, mode="fill", fill_value=0)
+    ref = jnp.take(ref, jnp.arange(3, 4), axis=2, mode="fill", fill_value=0)
+    ref = ref.squeeze(2)
+    assert jnp.all(jnp.equal(named1["H", ds(2, 5), "D", 3].array, ref))
+    ref = jnp.take(named1.array, jnp.arange(3, 8), axis=0, mode="fill", fill_value=0)
+    ref = jnp.take(ref, jnp.arange(3, 10, 2), axis=2, mode="fill", fill_value=0)
+    assert jnp.all(jnp.equal(named1["H", ds(3, 5), "D", 3:10:2].array, ref))
+
+
+def test_dslice_oob_read_and_write():
+    Seq = hax.Axis("seq", 5)
+    from haliax import ds
+
+    arr = hax.arange((Seq,), dtype=int)
+    out = arr[{"seq": ds(3, 4)}]
+    ref = jnp.take(arr.array, jnp.arange(3, 7), mode="fill", fill_value=0)
+    assert jnp.array_equal(out.array, ref)
+
+    upd = hax.arange((Seq.resize(4),), dtype=int)
+    updated = arr.at[{"seq": ds(3, 4)}].set(upd)
+    ref_upd = arr.array.at[jnp.arange(3, 7)].set(upd.array, mode="drop")
+    assert jnp.array_equal(updated.array, ref_upd)
 
 
 def test_slice_nd_array_present_dims():

{haliax-1.4.dev379 → haliax-1.4.dev381}/tests/test_ops.py

@@ -237,3 +237,16 @@ def test_reductions_produce_scalar_named_arrays_when_None_axis():
     # But if we specify axes, we always get a NamedArray, even if it's a scalar
     assert isinstance(hax.mean(named1, axis=("Height", "Width")), NamedArray)
     assert hax.mean(named1, axis=("Height", "Width")).axes == ()
+
+
+def test_pad():
+    Height = Axis("Height", 3)
+    Width = Axis("Width", 2)
+
+    arr = hax.arange((Height, Width))
+    padded = hax.pad(arr, {Height: (1, 2), Width: (0, 1)}, mode="constant", constant_values=0)
+
+    expected = jnp.pad(arr.array, [(1, 2), (0, 1)], mode="constant", constant_values=0)
+    assert padded.axes[0].size == Height.size + 3
+    assert padded.axes[1].size == Width.size + 1
+    assert jnp.all(expected == padded.array)

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

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