1

I have one tensor of the size of [batch_size,1] where each number for sample indicates an integer that is smaller than 5000. I'd like to create a new tensor of the size of [batch_size,5000] where the first numbers for each sample are ones, according to the first tensor, and the rest are zeros. For example:

t1=[[3],[5]]

so we will result with

t2=[[1,1,1,0,0,...],[1,1,1,1,1,0,0,...]]
Codevan
  • 538
  • 3
  • 20

1 Answers1

0

This is quite tricky since you want to fill values based on indices and not on the value itself...

Yet you can still manage it, but you have to get creative. We need some way for indices to be reflected on the values themselves. We will keep batch_size=2 and a vector size of 10:

>>> t1 = torch.LongTensor([[3], [5]])
>>> t2 = torch.zeros(len(t1), 10)
tensor([[0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]])

Here is the interesting part: arange a tensor with t2's shape and substract t1.

>>> torch.arange(10).repeat(2, 1) - t1
>>> tensor([[-3, -2, -1,  0,  1,  2,  3,  4,  5,  6],
            [-5, -4, -3, -2, -1,  0,  1,  2,  3,  4]])

Notice how values are negative before the breakpoint, and positive after. That's our mask:

>>> mask = torch.arange(10).repeat(2, 1) - t1 < 0
tensor([[ True,  True,  True, False, False, False, False, False, False, False],
        [ True,  True,  True,  True,  True, False, False, False, False, False]])  

To finish it off:

>>> t2[mask] = 1
tensor([[1., 1., 1., 0., 0., 0., 0., 0., 0., 0.],
        [1., 1., 1., 1., 1., 0., 0., 0., 0., 0.]])
Ivan
  • 34,531
  • 8
  • 55
  • 100