Make create_diagonal support broadcasting

Author: NeilGirdhar

src/array_api_extra/_lib/_funcs.py

@@ -5,7 +5,7 @@
 
 import math
 import warnings
-from collections.abc import Sequence
+from collections.abc import Generator, Sequence
 from types import ModuleType
 from typing import cast
 
@@ -163,6 +163,16 @@ def cov(m: Array, /, *, xp: ModuleType | None = None) -> Array:
     return xp.squeeze(c, axis=axes)
 
 
+def ndindex(*x: int) -> Generator[tuple[int, ...]]:
+    if not x:
+        yield ()
+        return
+    indices = list(ndindex(*x[1:]))
+    for i in range(x[0]):
+        for j in indices:
+            yield i, *j
+
+
 def create_diagonal(
     x: Array, /, *, offset: int = 0, xp: ModuleType | None = None
 ) -> Array:
@@ -172,7 +182,7 @@ def create_diagonal(
     Parameters
     ----------
     x : array
-        A 1-D array.
+        An array having shape (*broadcast_dims, k).
     offset : int, optional
         Offset from the leading diagonal (default is ``0``).
         Use positive ints for diagonals above the leading diagonal,
@@ -183,7 +193,8 @@ def create_diagonal(
     Returns
     -------
     array
-        A 2-D array with `x` on the diagonal (offset by `offset`).
+        An array having shape (*broadcast_dims, k+abs(offset), k+abs(offset)) with `x`
+        on the diagonal (offset by `offset`).
 
     Examples
     --------
@@ -206,18 +217,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))
+    pre = x.shape[:-1]
+    n = x.shape[-1] + abs(offset)
+    diag = xp.zeros((*pre, 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(*pre):
+        diag = at(diag)[(*index, target_slice)].set(x[*index, :])
+    return xp.reshape(diag, (*pre, n, n))
 
 
 def expand_dims(

tests/test_funcs.py

@@ -1,6 +1,7 @@
 import contextlib
 import warnings
 from types import ModuleType
+from math import prod
 
 import numpy as np
 import pytest
@@ -19,6 +20,7 @@
     sinc,
 )
 from array_api_extra._lib import Backend
+from array_api_extra._lib._funcs import ndindex
 from array_api_extra._lib._testing import xp_assert_close, xp_assert_equal
 from array_api_extra._lib._utils._compat import device as get_device
 from array_api_extra._lib._utils._typing import Array, Device
@@ -193,9 +195,17 @@ def test_0d(self, xp: ModuleType):
         with pytest.raises(ValueError, match="1-dimensional"):
             create_diagonal(xp.asarray(1))
 
+    @pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="no device kwarg in zeros()")
     def test_2d(self, xp: ModuleType):
-        with pytest.raises(ValueError, match="1-dimensional"):
-            create_diagonal(xp.asarray([[1]]))
+        result = create_diagonal(xp.asarray([[1]]))
+        xp_assert_equal(result, xp.asarray([[[1]]]))
+        b = xp.zeros((3, 2, 4, 5), dtype=xp.int64)
+        for i in ndindex(*b.shape):
+            b = at(b)[*i].set(prod(i))
+        c = create_diagonal(b)
+        zero = xp.zeros((), dtype=xp.int64)
+        for i in ndindex(*c.shape):
+            xp_assert_equal(c[*i], b[*(i[:-1])] if i[-2] == i[-1] else zero)
 
     @pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="no device kwarg in zeros()")
     def test_device(self, xp: ModuleType, device: Device):