Заполнение диагоналей квадратных матриц внутри трехмерного массива значениями, заданными двумерным массивом

Question

Для трехмерного ndarray z с формой (k,n,n) возможно ли без использования итерации заполнить диагонали матриц k nxn значениями, заданными 2D ndarray v, для формы (k,n)?

Например, результат операции должен быть таким же, как цикл по k матрицам:

z = np.zeros((3,10,10))
v = np.arange(30).reshape((3,10))

for i in range(len(z)): 
    np.fill_diagonal(z[i], v[i])

Есть ли способ сделать это без повторного вызова np.fill_diagonal внутри цикла? Если возможно, я бы предпочел решение, которое может быть применено и к массивам более высоких размеров, где z.shape == (a,b,c,...,k,n,n) и v.shape = (a,b,c,...,k,n)

Answer 1

Вот для общих n-dim массивов -

diag_view = np.einsum('...ii->...i',z)
diag_view[:] = v

Другой с изменением формы -

n = v.shape[-1] 
z.reshape(-1,n**2)[:,::n+1] = v.reshape(-1,n)
# or z.reshape(z.shape[:-2]+(-1,))[...,::n+1] = v

Другой с masking -

m = np.eye(n, dtype=bool) # n = v.shape[-1] from earlier
z[...,m] = v