Diagonal matrix for each batch entry of vector

Gluon
1

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.

2

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
3

Thank you very much @thomelane !!