Batch of diagonal matrix

Hi !
I have a matrix n*m of n different vectors of dimensions m.
I would like to get n matrices of size m*m with each matrix being a diagonal of a vector.

I guess I could do:

d = []
for vec in mat:
   d.append(torch.diag(vec))
torch.stack(d)

But isn’t there any better ‘pytorchic’ way ?

2 Likes

Yes! It’s kind of hidden, but here you go:

>>> mat = Variable(torch.randn(3, 4))
>>> res = Variable(torch.zeros(3, 4, 4))
>>> res.as_strided(mat.size(), [res.stride(0), res.size(2) + 1]).copy_(mat)
Variable containing:
 0.0422  0.0896  1.4919 -0.7167
-2.4854  0.3412 -1.4421 -0.5081
-0.6238  0.2446 -0.2848 -0.4184
[torch.FloatTensor of size (3,4)]

>>> res

(0 ,.,.) =
  0.0422  0.0000  0.0000  0.0000
  0.0000  0.0896  0.0000  0.0000
  0.0000  0.0000  1.4919  0.0000
  0.0000  0.0000  0.0000 -0.7167

(1 ,.,.) =
 -2.4854  0.0000  0.0000  0.0000
  0.0000  0.3412  0.0000  0.0000
  0.0000  0.0000 -1.4421  0.0000
  0.0000  0.0000  0.0000 -0.5081

(2 ,.,.) =
 -0.6238  0.0000  0.0000  0.0000
  0.0000  0.2446  0.0000  0.0000
  0.0000  0.0000 -0.2848  0.0000
  0.0000  0.0000  0.0000 -0.4184
[torch.FloatTensor of size 3x4x4]

I’ve created an issue for this at https://github.com/pytorch/pytorch/issues/5198

5 Likes

For the sake of completeness and to save future readers trouble, the function you are looking for is torch.diag_embed.

13 Likes

However, pytorch0.4.1 has not this function. How can I do this on pytorch0.4.1?