add diag_indices_from
@@ -167,3 +167,22 @@ def diag_indices(n, ndim=2):
"""
idx = cupy.arange(n)
return (idx,) * ndim
+
+
+def diag_indices_from(arr):
+ """
+ Return the indices to access the main diagonal of an n-dimensional array.
+ See `diag_indices` for full details.
+
+ Args:
+ arr (cupy.ndarray): At least 2-D.
+ """
+
+ if not arr.ndim >= 2:
+ raise ValueError("input array must be at least 2-d")
+ # For more than d=2, the strided formula is only valid for arrays with
+ # all dimensions equal, so we check first.
+ if not cupy.all(cupy.diff(arr.shape) == 0):
+ raise ValueError("All dimensions of input must be of equal length")
+
+ return diag_indices(arr.shape[0], arr.ndim)
