Merge branch 'main' into lazywhere

Author: crusaderky

.github/workflows/ci.yml

@@ -30,7 +30,7 @@ jobs:
       - uses: pre-commit/[email protected]
         with:
           extra_args: --hook-stage manual --all-files
-      - uses: prefix-dev/[email protected].2
+      - uses: prefix-dev/[email protected].3
         with:
           pixi-version: v0.40.3
           cache: true
@@ -56,7 +56,7 @@ jobs:
         with:
           fetch-depth: 0
 
-      - uses: prefix-dev/[email protected].2
+      - uses: prefix-dev/[email protected].3
         with:
           pixi-version: v0.40.3
           cache: true

.github/workflows/docs-build.yml

@@ -7,7 +7,7 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v4
-      - uses: prefix-dev/[email protected].2
+      - uses: prefix-dev/[email protected].3
         with:
           pixi-version: v0.40.3
           cache: true

pixi.lock

@@ -15,7 +15,7 @@
     sinc,
 )
 
-__version__ = "0.6.1.dev0"
+__version__ = "0.7.0.dev0"
 
 # pylint: disable=duplicate-code
 __all__ = [

src/array_api_extra/__init__.py

@@ -278,16 +278,11 @@ def _op(
             msg = f"copy must be True, False, or None; got {copy!r}"
             raise ValueError(msg)
 
-        if copy is None:
-            writeable = is_writeable_array(x)
-            copy = not writeable
-        elif copy:
-            writeable = None
-        else:
-            writeable = is_writeable_array(x)
+        writeable = None if copy else is_writeable_array(x)
 
-        # JAX inside jax.jit and Dask don't support in-place updates with boolean
-        # mask. However we can handle the common special case of 0-dimensional y
+        # JAX inside jax.jit doesn't support in-place updates with boolean
+        # masks; Dask exclusively supports __setitem__ but not iops.
+        # We can handle the common special case of 0-dimensional y
         # with where(idx, y, x) instead.
         if (
             (is_dask_array(idx) or is_jax_array(idx))
@@ -296,23 +291,24 @@ def _op(
         ):
             y_xp = xp.asarray(y, dtype=x.dtype)
             if y_xp.ndim == 0:
-                if out_of_place_op:
+                if out_of_place_op:  # add(), subtract(), ...
                     # suppress inf warnings on Dask
                     out = apply_where(
                         idx, (x, y_xp), out_of_place_op, fill_value=x, xp=xp
                     )
                     # Undo int->float promotion on JAX after _AtOp.DIVIDE
                     out = xp.astype(out, x.dtype, copy=False)
-                else:
+                else:  # set()
                     out = xp.where(idx, y_xp, x)
 
-                if copy:
-                    return out
-                x[()] = out
-                return x
+                if copy is False:
+                    x[()] = out
+                    return x
+                return out
+
             # else: this will work on eager JAX and crash on jax.jit and Dask
 
-        if copy:
+        if copy or (copy is None and not writeable):
             if is_jax_array(x):
                 # Use JAX's at[]
                 func = cast(Callable[[Array], Array], getattr(x.at[idx], at_op.value))
@@ -336,7 +332,7 @@ def _op(
             msg = f"Can't update read-only array {x}"
             raise ValueError(msg)
 
-        if in_place_op:
+        if in_place_op:  # add(), subtract(), ...
             x[self._idx] = in_place_op(x[self._idx], y)
         else:  # set()
             x[self._idx] = y

src/array_api_extra/_lib/_at.py

@@ -11,6 +11,8 @@
 
 from ._at import at
 from ._utils import _compat, _helpers
+from ._utils._compat import array_namespace, is_jax_array
+from ._utils._helpers import asarrays, ndindex
 from ._utils._compat import (
     array_namespace,
     is_dask_namespace,
@@ -384,7 +386,7 @@ def create_diagonal(
     Parameters
     ----------
     x : array
-        A 1-D array.
+        An array having shape ``(*batch_dims, k)``.
     offset : int, optional
         Offset from the leading diagonal (default is ``0``).
         Use positive ints for diagonals above the leading diagonal,
@@ -395,7 +397,8 @@ def create_diagonal(
     Returns
     -------
     array
-        A 2-D array with `x` on the diagonal (offset by `offset`).
+        An array having shape ``(*batch_dims, k+abs(offset), k+abs(offset))`` with `x`
+        on the diagonal (offset by `offset`).
 
     Examples
     --------
@@ -418,18 +421,21 @@ def create_diagonal(
     if xp is None:
         xp = array_namespace(x)
 
-    if x.ndim != 1:
-        err_msg = "`x` must be 1-dimensional."
+    if x.ndim == 0:
+        err_msg = "`x` must be at least 1-dimensional."
         raise ValueError(err_msg)
-    n = x.shape[0] + abs(offset)
-    diag = xp.zeros(n**2, dtype=x.dtype, device=_compat.device(x))
-
-    start = offset if offset >= 0 else abs(offset) * n
-    stop = min(n * (n - offset), diag.shape[0])
-    step = n + 1
-    diag = at(diag)[start:stop:step].set(x)
-
-    return xp.reshape(diag, (n, n))
+    batch_dims = x.shape[:-1]
+    n = x.shape[-1] + abs(offset)
+    diag = xp.zeros((*batch_dims, n**2), dtype=x.dtype, device=_compat.device(x))
+
+    target_slice = slice(
+        offset if offset >= 0 else abs(offset) * n,
+        min(n * (n - offset), diag.shape[-1]),
+        n + 1,
+    )
+    for index in ndindex(*batch_dims):
+        diag = at(diag)[(*index, target_slice)].set(x[(*index, slice(None))])
+    return xp.reshape(diag, (*batch_dims, n, n))
 
 
 def expand_dims(

src/array_api_extra/_lib/_funcs.py

@@ -3,6 +3,7 @@
 # https://github.com/scikit-learn/scikit-learn/pull/27910#issuecomment-2568023972
 from __future__ import annotations
 
+from collections.abc import Generator
 from types import ModuleType
 from typing import cast
 
@@ -182,6 +183,29 @@ def asarrays(
     return (xb, xa) if swap else (xa, xb)
 
 
+def ndindex(*x: int) -> Generator[tuple[int, ...]]:
+    """
+    Generate all N-dimensional indices for a given array shape.
+
+    Given the shape of an array, an ndindex instance iterates over the N-dimensional
+    index of the array. At each iteration a tuple of indices is returned, the last
+    dimension is iterated over first.
+
+    This has an identical API to numpy.ndindex.
+
+    Parameters
+    ----------
+    *x : int
+        The shape of the array.
+    """
+    if not x:
+        yield ()
+        return
+    for i in ndindex(*x[:-1]):
+        for j in range(x[-1]):
+            yield *i, j
+
+
 def meta_namespace(
     *arrays: Array | int | float | complex | bool | None,
     xp: ModuleType | None = None,