# Diagonal matrix for each batch entry of vector

Dear all,

I am breaking my head around this for some time, is there an efficient way (both nd and symbol) to create a diagonal matrix for each batch entry of a vector? Example, I have a matrix (dimensions BxN) `a = nd.random.uniform(shape=[5,3])`, where the first dimension corresponds to the batch dimension. I want to create another matrix, say b (shape -> [5,3,3]), such that for each index of the first dimension corresponds a diagonal matrix (3x3) with elements of a in the main diagonal.
Say:

``````a = nd.array([[0.1589, 0.6363, 0.3660],
[0.7480, 0.6879, 0.8742],
[0.4119, 0.9257, 0.5950],
[0.7859, 0.3367, 0.9770],
[0.5743, 0.7411, 0.2153]])
``````

then b is such that:

``````b[0] == [[0.1589, 0.0000, 0.0000],
[0.0000, 0.6363, 0.0000],
[0.0000, 0.0000, 0.3660]]

b[1] == [[0.7480, 0.0000, 0.0000],
[0.0000, 0.6879, 0.0000],
[0.0000, 0.0000, 0.8742]]
``````

and so on.

I can achieve this functionality with pytorch, using `torch.diag_embed(a)`, but I havenâ€™t managed to translate this in mxnet.

Thank you very much for your time.

Hi @feevos,

I can think of a way if you know N. Using NDArray this is simple to get using `.shape`, but youâ€™re probably looking for a hybridizable version, so youâ€™d need to provide N for the following method.

`F.broadcast_mul(a.expand_dims(2), F.eye(N))`

e.g. for NDArray:

`mx.nd.broadcast_mul(a.expand_dims(2), mx.nd.eye(N))`

Which gives an output of:

``````[[[0.1589 0.     0.    ]
[0.     0.6363 0.    ]
[0.     0.     0.366 ]]

[[0.748  0.     0.    ]
[0.     0.6879 0.    ]
[0.     0.     0.8742]]

[[0.4119 0.     0.    ]
[0.     0.9257 0.    ]
[0.     0.     0.595 ]]

[[0.7859 0.     0.    ]
[0.     0.3367 0.    ]
[0.     0.     0.977 ]]

[[0.5743 0.     0.    ]
[0.     0.7411 0.    ]
[0.     0.     0.2153]]]
<NDArray 5x3x3 @cpu(0)>
``````
1 Like

Thank you very much @thomelane !!