Filling tensor with zeros after certain index

nlp
1

Given a 3d tenzor, say: batch x sentence length x embedding dim

a = torch.rand((10, 1000, 96))

and an array(or tensor) of actual lengths for each sentence

lengths =  torch .randint(1000,(10,))

outputs tensor([ 370., 502., 652., 859., 545., 964., 566., 576.,1000., 803.])

How to fill tensor ‘a’ with zeros after certain index along dimension 1 (sentence length) according to tensor ‘lengths’ ?
I want smth like that :

a[ : , lengths : , : ]  = 0

One way of doing it (slow if batch size is big enough):

for i, length in enumerate( lengths ):
    a[ i , length  : , : ]  = 0
2

Hi,

Excuse me, if you have a tensor like this:

tensor([[[ 1.1937, -0.7235, -0.1802, -0.5610],
         [-0.7524,  1.3047, -1.5577,  1.8352],
         [ 1.1573, -1.8952,  0.4175, -0.2085]],

        [[ 0.4069, -0.1069, -0.3838, -0.2991],
         [-0.5824, -0.4965, -0.1542,  1.1482],
         [ 0.5182,  0.5445,  1.1730, -0.2523]]])

and lengths = torch.tensor([1, 2])
do you want to have sth like this?

tensor([[[ 1.1937, -0.7235, -0.1802, -0.5610],
         [ 0.0000,  0.0000,  0.0000,  0.0000],
         [ 1.1573, -1.8952,  0.4175, -0.2085]],

        [[ 0.4069, -0.1069, -0.3838, -0.2991],
         [-0.5824, -0.4965, -0.1542,  1.1482],
         [ 0.0000,  0.0000,  0.0000,  0.0000]]])

or sth like this:

tensor([[[ 1.4429, -0.8365, -0.2565,  0.2319],
         [ 0.0000,  0.0000,  0.0000,  0.0000],
         [ 0.0000,  0.0000,  0.0000,  0.0000]],

        [[ 1.7247,  0.2831,  0.0687,  0.6557],
         [ 0.5110,  1.7607,  0.8107,  0.3624],
         [ 0.0000,  0.0000,  0.0000,  0.0000]]])

Thanks

3

I want second option :

tensor([[[ 1.4429, -0.8365, -0.2565,  0.2319],
         [ 0.0000,  0.0000,  0.0000,  0.0000],
         [ 0.0000,  0.0000,  0.0000,  0.0000]],

        [[ 1.7247,  0.2831,  0.0687,  0.6557],
         [ 0.5110,  1.7607,  0.8107,  0.3624],
         [ 0.0000,  0.0000,  0.0000,  0.0000]]])
4

This works:

>>> import torch
>>> a = torch.rand(2, 3, 4)
>>> print(a)
tensor([[[0.5066, 0.5184, 0.1193, 0.6062],
         [0.4995, 0.1689, 0.6175, 0.7917],
         [0.7996, 0.0225, 0.1145, 0.4249]],

        [[0.0975, 0.2995, 0.5857, 0.0806],
         [0.4922, 0.4778, 0.9133, 0.1418],
         [0.6594, 0.4907, 0.3268, 0.1211]]])
>>> l = torch.randint(3, (2, ))
>>> print(l)
tensor([1, 2])
>>> l += torch.arange(2) * 3
>>> b = torch.arange(6).view(2, 3)
>>> b = b < l.view(2, 1)
>>> a * b.view(2, 3, 1).float()
tensor([[[0.5066, 0.5184, 0.1193, 0.6062],
         [0.0000, 0.0000, 0.0000, 0.0000],
         [0.0000, 0.0000, 0.0000, 0.0000]],

        [[0.0975, 0.2995, 0.5857, 0.0806],
         [0.4922, 0.4778, 0.9133, 0.1418],
         [0.0000, 0.0000, 0.0000, 0.0000]]])
>>>

I hope it helps !

2 Likes
5

What does this command do?

6

It creates a binary tensor whose values are calculated with respect to the comparison of l and b. The .view(2, 1) is for broadcasting purposes.

>>> import torch
>>> l = torch.randint(3, (2, ))
>>> l # the original form of the lengths
tensor([2, 1])
>>> l += torch.arange(2) * 3
>>> l # the modified form of the lengths for the comparison to work
tensor([2, 4])
>>> b = torch.arange(6).view(2, 3)
>>> b
tensor([[0, 1, 2],
        [3, 4, 5]])
>>> b < l.view(2, 1)
tensor([[1, 1, 0],
        [1, 0, 0]], dtype=torch.uint8)

The trick is to use torch.arange and then compare it to the indices you want to create the effect of zeroing starting from a given point. I hope it is clear !

1 Like
7

Thanks LeviViana
Nice solution
I also found a bit different approach here:

1 Like