Remove zeros from tensor but keeping dimensions

1

Suppose I have the following tensor:

a = torch.Tensor([[1,2,0],[0,-1,0]])

I want to remove zeros but keep the dimensions. The results should be:

[[1,2],[-1]]

Is there an easy way to achieve this?

2

You could try to use the experimental NestedTensor support as seen here:

a = torch.tensor([[1,2,0],[0,-1,0]])

a_list = [a_[a_!=0.] for a_ in a]
b = torch.nested.as_nested_tensor(a_list)
print(b)
# nested_tensor([
#   tensor([1, 2]),
#   tensor([-1])
# ])
1 Like
3

I had no idea this existed! Does it provide efficiency optimizations under the hood for things like batch-multiplications or is it just a semantic convenience? Seems like it’d be useful for dynamic sequence lengths instead of masking maybe?

4

Masking will produce the following output ([ 1., 2., -1.]) which is not the expected one.

import torch
a = torch.Tensor([[1,2,0],[0,-1,0]])

mask = (a != 0)

b = torch.masked_select(a, mask)
print(b)

1 Like
5

Oh, I was asking ptrblck about nested tensors.

6

I don’t know what exactly is supported using NestedTensors as this class is still in its beta/experimental state. The Transfomer layers are providing some nested tensor functionality as seen here and the fastpath uses it too as described in this blog post.

1 Like
7

Thanks @ptrblck for your reply. I was not aware such method exists. However, since it is beta, this is really suggesting that this approach is an anti-pattern approach. There should be better ways to handle my issue.

8

I don’t think there is any other way of handling tensors with different shapes besides the (beta) NestedTensor approach or just a plain Python list to store all objects.

9

No @ptrblck , I mean my issue where I was thinking to utilize this nested tensor approach. Luckily, I found other ways to deal with it.

Again, Thanks Much,.