Filling the diagonals of square matrices inside a 3D ndarray with values given by a 2D ndarray

Question

Given a 3D ndarray z with shape (k,n,n), is it possbile without using iteration to fill the diagonals of the k nxn matrices with values given by a

Answer 1

Here's for generic n-dim arrays -

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

Another with reshaping -

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

Another with masking -

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

Initializing output z

If we need to initialize the output array z and one that will cover for generic n-dim cases, it would be :

z = np.zeros(v.shape + (v.shape[-1],), dtype=v.dtype)

Then, we proceed with the earlier listed approaches.