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?

### >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
```