Make create_diagonal support broadcasting

Author: NeilGirdhar

src/array_api_extra/_lib/_funcs.py

@@ -9,6 +9,8 @@
 from types import ModuleType
 from typing import cast
 
+import numpy as np
+
 from ._at import at
 from ._utils import _compat, _helpers
 from ._utils._compat import array_namespace, is_jax_array
@@ -172,7 +174,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 +185,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 +209,20 @@ 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[0]),
+                         n + 1)
+    for index in np.ndindex(*pre):
+        indexed_x = x[*index, :]
+        diag = at(diag)[(*index, target_slice)].set(indexed_x)
+    return xp.reshape(diag, (*pre, n, n))
 
 
 def expand_dims(

tests/test_funcs.py

@@ -194,8 +194,8 @@ def test_0d(self, xp: ModuleType):
             create_diagonal(xp.asarray(1))
 
     def test_2d(self, xp: ModuleType):
-        with pytest.raises(ValueError, match="1-dimensional"):
-            create_diagonal(xp.asarray([[1]]))
+        result = create_diagonal(xp.asarray([[1]], dtype=xp.float64))
+        xp_assert_equal(result, xp.asarray([[[1]]], dtype=xp.float64))
 
     @pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="no device kwarg in zeros()")
     def test_device(self, xp: ModuleType, device: Device):