haliax 1.4.dev380__tar.gz → 1.4.dev382__tar.gz

{haliax-1.4.dev380 → haliax-1.4.dev382}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: haliax
-Version: 1.4.dev380
+Version: 1.4.dev382
 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/
@@ -33,14 +33,14 @@ Description-Content-Type: text/markdown
 <a href="">
 <img alt="License" src="https://img.shields.io/github/license/stanford-crfm/haliax?color=blue" />
 </a>
-<a href="https://https://pypi.org/project/haliax/">
+<a href="https://pypi.org/project/haliax/">
     <img alt="PyPI" src="https://img.shields.io/pypi/v/haliax?color=blue" />
 </a>
 
 > *Though you don’t seem to be much for listening, it’s best to be careful. If you managed to catch hold of even just a piece of my name, you’d have all manner of power over me.*<br/>
 > — Patrick Rothfuss, *The Name of the Wind*
 
-Haliax is a [JAX](https:://github.com/google/jax) library for building neural networks with named tensors, in the tradition of Alexander Rush's [Tensor Considered Harmful](https://nlp.seas.harvard.edu/NamedTensor).
+Haliax is a [JAX](https://github.com/google/jax) library for building neural networks with named tensors, in the tradition of Alexander Rush's [Tensor Considered Harmful](https://nlp.seas.harvard.edu/NamedTensor).
 Named tensors improve the **legibility** and **compositionality** of tensor programs by using named axes instead of positional indices
 as typically used in NumPy, PyTorch, etc.
 

{haliax-1.4.dev380 → haliax-1.4.dev382}/README.md

@@ -10,14 +10,14 @@
 <a href="">
 <img alt="License" src="https://img.shields.io/github/license/stanford-crfm/haliax?color=blue" />
 </a>
-<a href="https://https://pypi.org/project/haliax/">
+<a href="https://pypi.org/project/haliax/">
     <img alt="PyPI" src="https://img.shields.io/pypi/v/haliax?color=blue" />
 </a>
 
 > *Though you don’t seem to be much for listening, it’s best to be careful. If you managed to catch hold of even just a piece of my name, you’d have all manner of power over me.*<br/>
 > — Patrick Rothfuss, *The Name of the Wind*
 
-Haliax is a [JAX](https:://github.com/google/jax) library for building neural networks with named tensors, in the tradition of Alexander Rush's [Tensor Considered Harmful](https://nlp.seas.harvard.edu/NamedTensor).
+Haliax is a [JAX](https://github.com/google/jax) library for building neural networks with named tensors, in the tradition of Alexander Rush's [Tensor Considered Harmful](https://nlp.seas.harvard.edu/NamedTensor).
 Named tensors improve the **legibility** and **compositionality** of tensor programs by using named axes instead of positional indices
 as typically used in NumPy, PyTorch, etc.
 

{haliax-1.4.dev380 → haliax-1.4.dev382}/docs/indexing.md

@@ -1,7 +1,7 @@
 # Indexing and Slicing
 
 Haliax supports Numpy-style indexing, including so-called [Advanced Indexing](https://numpy.org/doc/stable/user/basics.indexing.html#advanced-indexing),
-though the syntax is necessarily different. Most forms of indexing are supporting, except we don't support indexing with
+though the syntax is necessarily different. Most forms of indexing are supported, except we don't support indexing with
 booleans right now. (JAX doesn't support indexing with non-constant bool arrays anyway,
 so I don't think it's worth the effort to implement it in Haliax.)
 
@@ -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.dev380 → haliax-1.4.dev382}/docs/nn.md

@@ -6,7 +6,7 @@
 Haliax provides a small number of neural network modules that are compatible with Equinox, though
 they naturally all use [haliax.NamedArray][]. (We welcome PRs for more modules! Nothing too exotic though.)
 
-The most interesting of these modules is [haliax.nn.Stacked][], which allows you to create homogenous "stacks"
+The most interesting of these modules is [haliax.nn.Stacked][], which allows you to create homogeneous "stacks"
 of the same module (e.g. transformer blocks), which is a common pattern in deep learning.
 
 ### Linear

{haliax-1.4.dev380 → haliax-1.4.dev382}/docs/state-dict.md

@@ -226,7 +226,7 @@ any Axis members to match the new shape.
 ::: haliax.state_dict.save_state_dict
 ::: haliax.state_dict.load_state_dict
 
-### Converting betweewn State Dicts and Modules
+### Converting between State Dicts and Modules
 
 ::: haliax.state_dict.from_state_dict
 ::: haliax.state_dict.to_state_dict

{haliax-1.4.dev380 → haliax-1.4.dev382}/docs/typing.md

@@ -1,4 +1,4 @@
-from haliax import NamedArrayfrom haliax import NamedArray
+from haliax import NamedArray
 
 # NamedArray Type Annotations
 

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

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

{haliax-1.4.dev380 → haliax-1.4.dev382}/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.dev380 → haliax-1.4.dev382}/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.dev380 → haliax-1.4.dev382}/src/haliax/nn/conv.py

@@ -213,7 +213,7 @@ class Conv(_ConvBase):
         return x
 
     def _do_conv(self, inputs):
-        # _do_conv expects there ot be a single __batch__ dimension
+        # _do_conv expects there to be a single __batch__ dimension
         output_axes = _compute_output_axes(inputs, "__batch__", self.In, self.Out)
 
         batch_index = _index_of_name(inputs.axes, "__batch__")

{haliax-1.4.dev380 → haliax-1.4.dev382}/src/haliax/nn/linear.py

@@ -143,7 +143,7 @@ class MoELinear(eqx.Module):
     Experts: AxisSpec = eqx.field(static=True)
     In: Axis = eqx.field(static=True)
     Out: Axis = eqx.field(static=True)
-    # TODO: support quanitization for ragged_dot?
+    # TODO: support quantization for ragged_dot?
     # dot_general: DotGeneralOp = eqx.field(default_factory=DotGeneralOp.default)
 
     use_gmm: bool = eqx.field(static=True)

{haliax-1.4.dev380 → haliax-1.4.dev382}/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.dev380 → haliax-1.4.dev382}/tests/test_rearrange.py

@@ -293,7 +293,7 @@ def test_examples():
     r = einops_rearrange(z, "{B (H: h1 h) (W: w1 w) C} -> (B: B h1 w1) ... (C: C h w) ", h1=2, w1=2)
     assert r.axes == (Axis("B", B.size * 2 * 2), D, Axis("C", C.size * sH.size * sW.size))
     # unet attention reordering:
-    # postional: (qkv heads c) h w -> qkv heads c (h w)
+    # positional: (qkv heads c) h w -> qkv heads c (h w)
     # named: { (embed: qkv heads c) h w } -> qkv heads c (pos: h w)
     Embed = Axis("embed", 3 * 4 * C.size)
     attn = hax.random.randint(PRNGKey(0), (Embed, H, W), 0, 255)

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

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