haliax 1.4.dev289__tar.gz → 1.4.dev290__tar.gz

{haliax-1.4.dev289 → haliax-1.4.dev290}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: haliax
-Version: 1.4.dev289
+Version: 1.4.dev290
 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.dev289 → haliax-1.4.dev290}/docs/fp8.md

@@ -1,7 +1,8 @@
 # FP8 Training
 
 !!! warning
-        FP8 training in Haliax is currently experimental and may change in the future.
+
+    FP8 training in Haliax is currently experimental and may change in the future.
 
 FP8 refers to 8-bit floating point numbers. FP8 is a massively reduced precision compared to the 32-bit floating point numbers
 or 16-bit floating point numbers that are typically used in deep learning: there are only 256 possible values in FP8, compared to
@@ -123,7 +124,7 @@ you will get the gradient computation as normal, but you'll also get the updated
 This updated state needs to directly replace the state in the module (rather than be used for a gradient step), which is
 why you need to use the `partition_for_grad_overwrite`
 
-The FP8 `dot_general` module is implemented in [haliax.quantization.Fp8DotGeneral][]. It's actually not that complicated:
+The FP8 `dot_general` module is implemented in [haliax.quantization.Fp8DotGeneralOp][]. It's actually not that complicated:
 
 1) It holds a scaling factor and history of maximum values for each of (lhs, rhs, output) and updates them based on the
 gradients.

{haliax-1.4.dev289 → haliax-1.4.dev290}/docs/indexing.md

@@ -37,12 +37,42 @@ Otherwise, the idea is pretty straightforward: any unspecified axes are treated
 slices are kept in reduced dimensions, and integers eliminate dimensions. If all dimensions are eliminated, a scalar
 JAX ndarray is returned.
 
+The following types are supported for indexing:
+
+* Integers, including scalar JAX arrays
+* Slices
+* [haliax.dslice][] objects (See [Dynamic Slices](#dynamic-slices) below.)
+* Lists of integers
+* Named arrays (See [Advanced Indexing](#advanced-indexing) below.)
+* 1-D JAX Arrays of integers
+
+1-D JAX Arrays are interpreted as NamedArrays with a single axis with the same name as
+the one they are slicing. That is:
+
+```python
+import haliax as hax
+import jax
+import jax.numpy as jnp
+
+X = hax.Axis("X", 10)
+Y = hax.Axis("Y", 20)
+
+a = hax.random.uniform(jax.random.PRNGKey(0), (X, Y))
+
+sliced = a["X", jnp.array([1, 2, 3])]
+
+# same as
+a.array[jnp.array([1, 2, 3]), :]
+```
+
+Note that boolean arrays are not supported, as JAX does not support them in JIT-compiled code. You
+can use [haliax.where][] for most of the same functionality, though.
 
 ### Shapes in JAX
 
-Before we continue note on shapes in JAX. Most JAX code will be used inside `jit`, which means that the sizes of all arrays
-must be determined at compile time (i.e. when JAX interprets your functions abstractly). This is a hard requirement in
-XLA.
+Before we continue, a note on shapes in JAX. Most JAX code will be used inside `jit`, which means that the sizes of all
+arrays must be determined at compile time (i.e. when JAX interprets your functions abstractly). This is a hard
+requirement in XLA. It might worked around one day, but it's the way things are for now.
 
 A consequence of this restriction is that certain indexing patterns aren't allowed in `jit`-ed JAX code:
 
@@ -177,3 +207,60 @@ a[{"Y": ind1}]  # error, "X" is not eliminated by the indexing operation
 
 a[{"X": ind2, "Y": ind1}]  # ok, because X and Y are eliminated by the indexing operation
 ```
+
+## Index Update
+
+JAX is a functional version of NumPy, so it doesn't directly support in-place updates. It does
+however [provide an `at` syntax](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.ndarray.at.html#jax.numpy.ndarray.at)
+to express the same logic (and that will typically be optimized to be as efficient as an in-place update). Haliax
+provides a similar syntax for updating arrays.
+
+```python
+import haliax as hax
+
+X = hax.Axis("X", 10)
+Y = hax.Axis("Y", 20)
+Z = hax.Axis("Z", 30)
+
+a = hax.zeros((X, Y, Z))
+
+a.at[{"X": 1, "Y": 2, "Z": 3}].set(1.0)  # sets a[1, 2, 3] to 1.0
+a.at["X", 1].set(2.0)  # sets a[1, :, :] to 2.0
+
+a.at[{"X": 1, "Y": hax.ds(3, 5), "Z": 3}].add(1.0)  # adds 1.0 to a[1, 3:8, 3]
+```
+
+Haliax supports the same `at` functionality as JAX, just with named arrays and additionally dslices. A summary of the
+`at` syntax is as follows:
+
+| Alternate Syntax             | Equivalent In-Place Operation |
+|------------------------------|-------------------------------|
+| `x = x.at[idx].set(y)`       | `x[idx] = y`                  |
+| `x = x.at[idx].add(y)`       | `x[idx] += y`                 |
+| `x = x.at[idx].multiply(y)`  | `x[idx] *= y`                 |
+| `x = x.at[idx].divide(y)`    | `x[idx] /= y`                 |
+| `x = x.at[idx].power(y)`     | `x[idx] **= y`                |
+| `x = x.at[idx].min(y)`       | `x[idx] = minimum(x[idx], y)` |
+| `x = x.at[idx].max(y)`       | `x[idx] = maximum(x[idx], y)` |
+| `x = x.at[idx].apply(ufunc)` | `ufunc.at(x, idx)`            |
+| `x = x.at[idx].get()`        | `x = x[idx]`                  |
+
+These methods also have options to control out-of-bounds behavior, as well as allowing you
+to specify that the indices are sorted or unique. (If they are, XLA can sometimes optimize the
+operation more effectively.)
+
+!!! note
+
+    These named arguments are not passed to `at`, but to the next method in the chain.
+
+(This is copied from the JAX documentation:)
+
+* `mode`: One of `"promise_in_bounds"`, `"clip"`, `"drop"`, or `"fill"`. See [JAX's documentation](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.GatherScatterMode.html#jax.lax.GatherScatterMode) for more details.
+* `indices_are_sorted`: If `True`, the implementation will assume that the indices passed to `at` are sorted in ascending order, which can lead to more efficient execution on some backends.
+* `unique_indices`: If `True`, the implementation will assume that the indices passed to `at` are unique, which can result in more efficient execution on some backends.
+* `fill_value`: Only applies to the `get()` method: the fill value to return for out-of-bounds slices when mode is 'fill'. Ignored otherwise. Defaults to NaN for inexact types, the largest negative value for signed types, the largest positive value for unsigned types, and True for booleans.
+
+!!! tip
+
+    It's worth emphasizing that these functions are typically compiled to scatter-add and friends (as appropriate).
+    This is the preferred way to do scatter/gather operations in JAX, as well as in Haliax.

{haliax-1.4.dev289 → haliax-1.4.dev290}/docs/rearrange.md

@@ -54,6 +54,7 @@ If you're used to einops, the syntax should be familiar, with the main differenc
 and the additional "unordered" syntax for selecting dimensions by name.
 
 !!! warning
+
     This syntax is fairly new. It is pretty well-tested, but it is possible that there are bugs.
 
 ### Examples

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

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

{haliax-1.4.dev289 → haliax-1.4.dev290}/src/haliax/core.py

@@ -9,7 +9,6 @@ from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, Tuple
 
 import jax
 import jax.numpy as jnp
-import numpy
 import numpy as np
 
 import haliax
@@ -18,22 +17,21 @@ from haliax.jax_utils import is_jax_array_like, is_pallas_dslice
 from haliax.util import ensure_tuple
 
 from ._src.util import index_where, py_slice, slice_t
-from .axis import (
-    Axis,
-    AxisSelection,
-    AxisSelector,
-    AxisSpec,
-    axis_name,
-    dslice,
-    eliminate_axes,
-    selects_axis,
-    union_axes,
-)
-from .types import DTypeLike, IntScalar, PrecisionLike, Scalar
+from .axis import Axis, AxisSelection, AxisSelector, AxisSpec, axis_name, dslice, eliminate_axes, selects_axis
+from .types import GatherScatterModeStr, IntScalar, PrecisionLike, Scalar
 
 
 NamedOrNumeric = Union[Scalar, "NamedArray"]
-NamedIndex = Union[int, slice_t, "NamedArray", dslice]
+NamedIndex = Union[int, slice_t, "NamedArray", dslice, list[int], jnp.ndarray]
+
+SliceSpec = Union[
+    tuple[AxisSelector, NamedIndex],
+    tuple[AxisSelector, NamedIndex, AxisSelector, NamedIndex],
+    tuple[AxisSelector, NamedIndex, AxisSelector, NamedIndex, AxisSelector, NamedIndex],
+    tuple[AxisSelector | NamedOrNumeric, ...],
+    Mapping[AxisSelector, NamedIndex],
+]
+
 
 _ENABLE_SHAPE_CHECKS = True
 
@@ -330,19 +328,47 @@ class NamedArray:
     def take(self, axis: AxisSelector, index: Union[int, "NamedArray"]) -> "NamedArray":
         return haliax.take(self, axis=axis, index=index)
 
-    @overload
-    def __getitem__(self, item: Tuple[AxisSelector, NamedIndex]) -> "NamedArray":
-        ...
+    @property
+    def at(self) -> "_NamedIndexUpdateHelper":
+        """
+                Named analog of [jax's at method](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.ndarray.at.html).
 
-    @overload
-    def __getitem__(self, item: Tuple[AxisSelector, NamedIndex, AxisSelector, NamedIndex]) -> "NamedArray":
-        ...
+                Docs from the JAX docs:
 
-    @overload
-    def __getitem__(self, item: Mapping[AxisSelector, NamedIndex]) -> "NamedArray":
-        ...
+                The at property provides a functionally pure equivalent of in-place array modifications.
+
+                In particular:
+
+        | Alternate syntax | Equivalent In-place expression |
+        |------------------|-------------------------------|
+        | `x = x.at[idx].set(y)` | `x[idx] = y` |
+        | `x[idx] = y` | `x = x.at[idx].set(y)` |
+        | `x = x.at[idx].add(y)` | `x[idx] += y`|
+        | `x = x.at[idx].multiply(y)` | `x[idx] *= y`|
+        | `x = x.at[idx].divide(y)` | `x[idx] /= y`|
+        | `x = x.at[idx].power(y)` | `x[idx] **= y`|
+        | `x = x.at[idx].min(y)` | `x[idx] = minimum(x[idx], y)`|
+        | `x = x.at[idx].max(y)` | `x[idx] = maximum(x[idx], y)`|
+        | `x = x.at[idx].apply(ufunc)` | `ufunc.at(x, idx)`|
+
+        x = x.at[idx].get()
+
+
+        x = x[idx]
+
+        None of the x.at expressions modify the original x; instead they return a modified copy of x. However, inside a jit() compiled function, expressions like x = x.at[idx].set(y) are guaranteed to be applied in-place.
+
+        Unlike NumPy in-place operations such as x[idx] += y, if multiple indices refer to the same location, all updates will be applied (NumPy would only apply the last update, rather than applying all updates.) The order in which conflicting updates are applied is implementation-defined and may be nondeterministic (e.g., due to concurrency on some hardware platforms).
+
+        By default, JAX assumes that all indices are in-bounds. Alternative out-of-bound index semantics can be specified via the mode parameter (see below).
+
+
+                Returns:
 
-    def __getitem__(self, idx) -> "NamedArray":
+        """
+        return _NamedIndexUpdateHelper(self)
+
+    def __getitem__(self, idx: SliceSpec) -> "NamedArray":
         """Syntactic sugar for [haliax.index][], which is the actual implementation.
 
         Supports indexing like:
@@ -366,17 +392,8 @@ class NamedArray:
 
         This returns a NamedArray if any axes remain, or a scalar (0-dimensional) jnp.ndarray if all axes are indexed out.
         """
-        if isinstance(idx, tuple):
-            if len(idx) == 1:
-                idx = idx[0]
-            else:
-                if len(idx) % 2 != 0:
-                    raise ValueError(
-                        "Must provide an even number of arguments to __getitem__ when using the shorthand syntax."
-                    )
-                idx = {idx[i]: idx[i + 1] for i in range(0, len(idx), 2)}
-
-        return index(self, idx)
+        idx_dict = _convert_index_expr_to_dict(idx)
+        return index(self, idx_dict)
 
     # np.ndarray methods:
     def all(self, axis: Optional[AxisSelection] = None, *, where: Optional["NamedArray"] = None) -> "NamedArray":
@@ -886,28 +903,60 @@ 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.
 
+    See Also:
+        * [haliax.NamedArray.at][] for a functional equivalent of in-place array modifications.
+
     Returns:
         NamedArray or jnp.ndarray: A NamedArray is returned if there are any axes remaining after selection,
         otherwise a scalar (0-dimensional) jnp.ndarray is returned if all axes are indexed out.
     """
     # indices where we have array args
-    array_slice_indices = []
-    dslice_indices = []
+    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)]
+
+    return haliax.named(sliced, new_axes)
+
+
+def _compute_new_axes_and_slices_for_index(
+    array, slices
+) -> tuple[AxisSpec, list[py_slice | dslice | jnp.ndarray | int | list[int]]]:
     ordered_slices: list = [py_slice(None, None, None)] * len(array.axes)  # type: ignore
     kept_axes = [True] * len(array.axes)
+    array_slice_indices = []
+
     for axis, slice_ in slices.items():
         axis_index = array._lookup_indices(axis)
         if axis_index is None:
             raise ValueError(f"axis {axis} not found in {array}")
-        ordered_slices[axis_index] = slice_
-
-        kept_axes[axis_index] = isinstance(slice_, py_slice) or isinstance(slice_, dslice) or is_pallas_dslice(slice_)
-
-        if isinstance(slice_, NamedArray):
+        if isinstance(slice_, py_slice) or isinstance(slice_, dslice) or is_pallas_dslice(slice_):
+            ordered_slices[axis_index] = slice_
+            kept_axes[axis_index] = True
+        elif isinstance(slice_, int):
+            ordered_slices[axis_index] = slice_
+            kept_axes[axis_index] = False
+        elif isinstance(slice_, NamedArray):
+            ordered_slices[axis_index] = slice_
             array_slice_indices.append(axis_index)
-
-        if isinstance(slice_, dslice) or is_pallas_dslice(slice_):
-            dslice_indices.append(axis_index)
+            kept_axes[axis_index] = False
+        elif isinstance(slice_, list):
+            # we'll let JAX complain if this is wrong
+            ordered_slices[axis_index] = slice_
+        elif isinstance(slice_, jnp.ndarray):
+            # we allow this if it's a 0-d or 1-d array
+            if slice_.ndim == 0:
+                ordered_slices[axis_index] = slice_
+            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))
+                kept_axes[axis_index] = False
+                array_slice_indices.append(axis_index)
+            else:
+                raise ValueError(
+                    f"Only 0-d or 1-d unnamed arrays can be used for indexing. Got {slice_} for axis {axis}"
+                )
+        else:
+            raise ValueError(f"Only NamedArrays can be used for advanced indexing. Got {slice_} for axis {axis}")
 
     # advanced indexing
     if len(array_slice_indices) > 0:
@@ -958,23 +1007,39 @@ def index(array: NamedArray, slices: Mapping[AxisSelector, NamedIndex]) -> Named
             # the advanced indices are not contiguous, so we need to insert the new axes at the front
             new_axes = broadcasted_axes + tuple(ax for i, ax in enumerate(array.axes) if kept_axes[i])
     else:
-        new_axes = tuple(axis.name for axis, keep in zip(array.axes, kept_axes) if keep)
+        new_axes = tuple(axis for axis, keep in zip(array.axes, kept_axes) if keep)
 
-    sliced = array.array
+    new_axes = tuple(axis_name(ax) for ax in new_axes)
+    return new_axes, ordered_slices
 
-    if len(dslice_indices) > 0:
-        # dynamic slice out the dslices
-        indices = [0] * len(array.axes)
-        lengths = [ax.size for ax in array.axes]
-        for i in dslice_indices:
-            indices[i] = ordered_slices[i].start
-            lengths[i] = ordered_slices[i].size
-        sliced = jax.lax.dynamic_slice(sliced, indices, lengths)
-        for i in dslice_indices:
-            ordered_slices[i] = py_slice(None, None, None)
 
-    sliced = sliced[tuple(ordered_slices)]
-    return haliax.named(sliced, new_axes)
+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]:
@@ -1445,6 +1510,208 @@ def flatten_all_axes_but(
     return result, unflatten
 
 
+class _NamedIndexUpdateHelper:
+    def __init__(self, array: NamedArray):
+        self.array = array
+
+    def __getitem__(self, slices: SliceSpec) -> "_NamedIndexUpdateRef":
+        return _NamedIndexUpdateRef(self.array, _convert_index_expr_to_dict(slices))
+
+
+class _NamedIndexUpdateRef:
+    def __init__(self, array: NamedArray, slices: SliceSpec):
+        self._array = array
+        self._slices = slices
+
+    def get(
+        self,
+        *,
+        indices_are_sorted: bool = False,
+        unique_indices: bool = False,
+        mode: Optional[GatherScatterModeStr] = None,
+        fill_value: Optional[Scalar] = None,
+    ) -> NamedArray:
+        slices, sliced_axes = _raw_indices_for_at(self._array, self._slices)
+        new_array = self._array.array.at[tuple(slices)].get(
+            indices_are_sorted=indices_are_sorted, unique_indices=unique_indices, mode=mode, fill_value=fill_value
+        )
+        return NamedArray(new_array, sliced_axes)
+
+    def set(
+        self,
+        update: NamedOrNumeric,
+        *,
+        indices_are_sorted: bool = False,
+        unique_indices: bool = False,
+        mode: Optional[GatherScatterModeStr] = None,
+    ) -> NamedArray:
+        slices, sliced_axes = _raw_indices_for_at(self._array, self._slices)
+        update = haliax.broadcast_to(update, sliced_axes, enforce_no_extra_axes=True)
+        new_array = self._array.array.at[tuple(slices)].set(
+            update.array, indices_are_sorted=indices_are_sorted, unique_indices=unique_indices, mode=mode
+        )
+        return NamedArray(new_array, self._array.axes)
+
+    def add(
+        self,
+        update: NamedOrNumeric,
+        *,
+        indices_are_sorted: bool = False,
+        unique_indices: bool = False,
+        mode: Optional[GatherScatterModeStr] = None,
+    ) -> NamedArray:
+        slices, sliced_axes = _raw_indices_for_at(self._array, self._slices)
+        update = haliax.broadcast_to(update, sliced_axes, enforce_no_extra_axes=True)
+        new_array = self._array.array.at[tuple(slices)].add(
+            update.array, indices_are_sorted=indices_are_sorted, unique_indices=unique_indices, mode=mode
+        )
+        return NamedArray(new_array, self._array.axes)
+
+    def multiply(
+        self,
+        update: NamedOrNumeric,
+        *,
+        indices_are_sorted: bool = False,
+        unique_indices: bool = False,
+        mode: Optional[GatherScatterModeStr] = None,
+    ) -> NamedArray:
+        slices, sliced_axes = _raw_indices_for_at(self._array, self._slices)
+        update = haliax.broadcast_to(update, sliced_axes, enforce_no_extra_axes=True)
+        new_array = self._array.array.at[tuple(slices)].multiply(
+            update.array, indices_are_sorted=indices_are_sorted, unique_indices=unique_indices, mode=mode
+        )
+        return NamedArray(new_array, self._array.axes)
+
+    def divide(
+        self,
+        update: NamedOrNumeric,
+        *,
+        indices_are_sorted: bool = False,
+        unique_indices: bool = False,
+        mode: Optional[GatherScatterModeStr] = None,
+    ) -> NamedArray:
+        slices, sliced_axes = _raw_indices_for_at(self._array, self._slices)
+        update = haliax.broadcast_to(update, sliced_axes, enforce_no_extra_axes=True)
+        new_array = self._array.array.at[tuple(slices)].divide(
+            update.array, indices_are_sorted=indices_are_sorted, unique_indices=unique_indices, mode=mode
+        )
+        return NamedArray(new_array, self._array.axes)
+
+    def max(
+        self,
+        update: NamedOrNumeric,
+        *,
+        indices_are_sorted: bool = False,
+        unique_indices: bool = False,
+        mode: Optional[GatherScatterModeStr] = None,
+    ) -> NamedArray:
+        slices, sliced_axes = _raw_indices_for_at(self._array, self._slices)
+        update = haliax.broadcast_to(update, sliced_axes, enforce_no_extra_axes=True)
+        new_array = self._array.array.at[tuple(slices)].max(
+            update.array, indices_are_sorted=indices_are_sorted, unique_indices=unique_indices, mode=mode
+        )
+        return NamedArray(new_array, self._array.axes)
+
+    def min(
+        self,
+        update: NamedOrNumeric,
+        *,
+        indices_are_sorted: bool = False,
+        unique_indices: bool = False,
+        mode: Optional[GatherScatterModeStr] = None,
+    ) -> NamedArray:
+        slices, sliced_axes = _raw_indices_for_at(self._array, self._slices)
+        update = haliax.broadcast_to(update, sliced_axes, enforce_no_extra_axes=True)
+        new_array = self._array.array.at[tuple(slices)].min(
+            update.array, indices_are_sorted=indices_are_sorted, unique_indices=unique_indices, mode=mode
+        )
+        return NamedArray(new_array, self._array.axes)
+
+    def power(
+        self,
+        update: NamedOrNumeric,
+        *,
+        indices_are_sorted: bool = False,
+        unique_indices: bool = False,
+        mode: Optional[GatherScatterModeStr] = None,
+    ) -> NamedArray:
+        slices, sliced_axes = _raw_indices_for_at(self._array, self._slices)
+        update = haliax.broadcast_to(update, sliced_axes, enforce_no_extra_axes=True)
+        new_array = self._array.array.at[tuple(slices)].power(
+            update.array, indices_are_sorted=indices_are_sorted, unique_indices=unique_indices, mode=mode
+        )
+        return NamedArray(new_array, self._array.axes)
+
+    def apply(
+        self,
+        func,
+        *,
+        indices_are_sorted: bool = False,
+        unique_indices: bool = False,
+        mode: Optional[GatherScatterModeStr] = None,
+    ) -> NamedArray:
+        # It's not really documented, but func can be any callable that takes a scalar array and returns a scalar array
+        slices, sliced_axes = _raw_indices_for_at(self._array, self._slices)
+        new_array = self._array.array.at[tuple(slices)].apply(
+            func, indices_are_sorted=indices_are_sorted, unique_indices=unique_indices, mode=mode
+        )
+        return NamedArray(new_array, self._array.axes)
+
+
+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
+    return ordered_slices, _sliced.axes
+
+
+def _convert_index_expr_to_dict(idx) -> dict[AxisSelector, NamedIndex]:
+    if isinstance(idx, tuple):
+        if len(idx) == 1:
+            idx = idx[0]
+        else:
+            if len(idx) % 2 != 0:
+                raise ValueError(
+                    "Must provide an even number of arguments to __getitem__ when using the shorthand syntax."
+                )
+            idx = {idx[i]: idx[i + 1] for i in range(0, len(idx), 2)}
+    elif isinstance(idx, dict):
+        pass
+    else:
+        raise ValueError(f"Invalid index type {type(idx)}")
+    return idx
+
+
 __all__ = [
     "NamedArray",
     "named",

{haliax-1.4.dev289 → haliax-1.4.dev290}/src/haliax/types.py

@@ -1,4 +1,4 @@
-from typing import Any, Protocol, Tuple, TypeAlias, Union
+from typing import Any, Literal, Protocol, Tuple, TypeAlias, Union
 
 import jax.numpy as jnp
 import numpy as np
@@ -28,3 +28,5 @@ Scalar = Union[float, int, jnp.ndarray]  # ndarray b/c array(1) is a scalar
 IntScalar = Union[int, jnp.ndarray]
 
 PrecisionLike = Union[None, str, Precision, Tuple[str, str], Tuple[Precision, Precision]]
+
+GatherScatterModeStr = Literal["promise_in_bounds", "clip", "drop", "fill"]

{haliax-1.4.dev289 → haliax-1.4.dev290}/tests/core_test.py

@@ -649,6 +649,34 @@ def test_slice_nd_array_present_dims():
     assert jnp.all(jnp.equal(named1[{"H": index2}].array, named1.array[index2.array, :, :]))
 
 
+def test_slice_nd_array_unnamed_slice():
+    # tests slicing with arrays that are already present in the named array, which is sometimes ok
+    H = Axis("H", 10)
+    W = Axis("W", 20)
+    D = Axis("D", 30)
+
+    named1 = hax.random.uniform(PRNGKey(0), (H, W, D))
+
+    index1 = jax.random.randint(PRNGKey(1), (4,), 0, H.size)
+    assert jnp.all(jnp.equal(named1[{"H": index1}].array, named1.array[index1, :, :]))
+
+    # hidden behavior: if we also pass in an H index to e.g. D, it is zipped together
+    index2 = hax.random.randint(PRNGKey(2), Axis("H", 4), 0, D.size)
+    assert jnp.all(jnp.equal(named1[{"H": index1, "D": index2}].array, named1.array[index1, :, index2.array]))
+
+    # this is different though:
+    index2r = index2.array
+    assert jnp.all(
+        jnp.equal(
+            named1[{"H": index1, "D": index2r}].array, named1.array[index1.reshape(1, -1), :, index2r.reshape(-1, 1)]
+        )
+    )
+    assert named1[{"H": index1, "D": index2r}].shape != named1[{"H": index1, "D": index2}].shape
+
+    index1 = list(index1)
+    assert jnp.all(jnp.equal(named1[{"H": index1}].array, named1.array[index1, :, :]))
+
+
 def test_full_indexing_returns_named_array():
     H = Axis("H", 10)
     W = Axis("W", 20)
@@ -847,3 +875,82 @@ def test_named_arrays_work_in_eqxi_while_loop():
     grad_fun = eqx.filter_value_and_grad(loss_fun)
 
     grad_fun(named1)
+
+
+def test_at_for_in_placeish():
+    H = Axis("H", 10)
+    W = Axis("W", 20)
+
+    named1 = hax.random.uniform(PRNGKey(0), (H, W))
+
+    named1_at = named1.at[H, 0].set(0)
+
+    assert jnp.all(jnp.equal(named1_at[H, 0].array, 0))
+    assert jnp.all(named1_at[H, 1:].array == named1[H, 1:].array)
+
+    # test add, multiply, power, etc.
+    named1_at = named1.at[H, 0].add(1)
+    assert jnp.all(named1_at.array == named1.array.at[0].add(1))
+
+    named1_at = named1.at[H, 0].multiply(2)
+    assert jnp.all(named1_at.array == named1.array.at[0].multiply(2))
+
+    named1_at = named1.at[H, 0].power(2)
+    assert jnp.all(named1_at.array == named1.array.at[0].power(2))
+
+    named1_at = named1.at[H, 0].divide(2)
+    assert jnp.all(named1_at.array == named1.array.at[0].divide(2))
+
+    named1_at = named1.at[H, 0].apply(hax.square)
+    assert jnp.all(named1_at.array == named1.array.at[0].apply(jnp.square))
+
+    named1_at = named1.at[H, 0].max(0.5)
+    assert jnp.all(named1_at.array == named1.array.at[0].max(0.5))
+
+    named1_at = named1.at[H, 0].min(0.5)
+    assert jnp.all(named1_at.array == named1.array.at[0].min(0.5))
+
+
+def test_at_with_fancy_indexing():
+    H = Axis("H", 10)
+    W = Axis("W", 20)
+    I0 = Axis("I0", 5)
+    I1 = Axis("I1", 5)
+
+    named1 = hax.random.uniform(PRNGKey(0), (H, W))
+    ind1 = hax.random.randint(PRNGKey(0), (I0,), 0, H.size)
+    ind2 = hax.random.randint(PRNGKey(0), (I1,), 0, W.size)
+
+    named1_at = named1.at[H, ind1].set(0)
+    assert jnp.all(named1_at.array == named1.array.at[ind1.array].set(0))
+
+    named1_at = named1.at[H, ind1].add(1, mode="clip")
+    assert jnp.all(named1_at.array == named1.array.at[ind1.array].add(1, mode="clip"))
+
+    named1_at = named1.at[H, ind1, W, ind2].set(0)
+    assert jnp.all(named1_at.array == named1.array.at[ind1.array.reshape(-1, 1), ind2.array.reshape(1, -1)].set(0))
+
+    # dslices
+    from haliax import ds
+
+    named1_at = named1.at[H, ds(3, 5)].set(0)
+    assert jnp.all(named1_at.array == named1.array.at[3:8].set(0))
+
+    named1_at = named1.at[H, ds(3, 5), W, ind2].power(2)
+    assert jnp.all(named1_at.array == named1.array.at[3:8, ind2.array].power(2))
+
+
+def test_slice_dslice_and_array():
+    H = Axis("H", 10)
+    W = Axis("W", 20)
+    I0 = Axis("I0", 5)
+
+    named1 = hax.random.uniform(PRNGKey(0), (H, W))
+    ind2 = hax.random.randint(PRNGKey(0), (I0,), 0, W.size)
+
+    from haliax import ds
+
+    named1.array.at[3:8, ind2.array].add(jnp.full((5, 5), 2))
+
+    named1_at = named1.at[H, ds(3, 5), W, ind2].add(2)
+    assert jnp.all(named1_at.array == named1.array.at[3:8, ind2.array].add(2))

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

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