Masked_select -- inverse operation?

1

Hello,

The masked_select function allows me to extract elements from a tensor. Given a tensor A and a boolean mask mask, masked_select will give me the indexed elements as a 1D-vector (call it b), such that

b = torch.masked_select(A, mask)

My question is if there is an efficient way would be to achieve the inverse operation? This is, given both b and mask, how can I recover A or rather a version of A where the indexed elements are in the right position and all-non indexed elements of A can be some constant, e.g. zero.

My use case is that A is about 15k x 15k dimensions and b holds about `200k elements. So I need something fast that should also work with auto-grad. When I am implementing a simple for loop, it takes ages to propagate back through the graph.

1 Like
2

You could try it with this code sample:

A = torch.randn(10, 10)
mask = torch.ByteTensor(10, 10).random_(2)
b = torch.masked_select(A, mask)

C = torch.zeros(10, 10)
C[mask] = b
print(A, mask, C)

I’m not sure if it fits your needs regarding the functionality with auto-grad.
Could you test it or give a small code sample showing what you would like to do?

5 Likes
3

Hey, this works a treat! Thank you. It’s appropriate for my use case, the speed up is huge. It does work with autograd too. Thanks a lot!

4

Note, the inverse operation is called masked_scatter.