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

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])

MEDevel.com: Open-source for Healthcare and Education

Collecting and validating open-source software for healthcare, education, enterprise, development, medical imaging, medical records, and digital pathology.

Visit Medevel

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?

>Solution :

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