Expand a dimension of 3-dimensional array into a diagonal matrix with vectorized computations-CodePudding

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?

CodePudding user response：

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