I have np.ndarray A of shape (N, M, D).
I’d like to create np.ndarray B of shape (N, M, D, D) such that for every pair of fixed indices n, m along axes 0 and 1
B[n, m] = np.eye(A[n, m])
I understand how to solve this problem using cycles, yet I’d like to write code performing this in vectorized manner. How can this be done using numpy?
import numpy as np A = ... # Your array here n, m, d = A.shape indices = np.arange(d) B = np.zeros((n, m, d, d)) B[:, :, indices, indices] = A