# Filling tensor with zeros after certain index

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``````

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

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]]])
``````

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

What does this command do?

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

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

1 Like