Conditionally apply tensor operations in PyTorch

Question

I know PyTorch doesn't have a map-like function to apply a function to each element of a tensor. So, could I do something like the following without a map-like function in PyTorch?

if tensor_a * tensor_b.matmul(tensor_c) < 1:
    return -tensor_a*tensor_b
else:
    return 0

This would work if the tensors were 1D. However, I need this to work when tensor_b is 2D (tensor_a needs to be unsqueezed in the return statement). This means a 2D tensor should be returned where some of the rows will be 0 vectors.

Happy to use the latest features of the most recent Python version.

Answer 1

If I understand correctly, you are looking to return a tensor either way (hence the mapping) but by checking the condition element-wise. Assuming the shapes of tensor_a, tensor_b, and tensor_c are all two dimensional, as in "simple matrices", here is a possible solution.

What you're looking for is probably torch.where, it's fairly close to a mapping where based on a condition, it will return one value or another element-wise.

It works like torch.where(condition, value_if, value_else) where all three tensors have the same shape (value_if and value_else can actually be floats which will be cast to tensors, filled with the same value). Also, condition is a bool tensor which defines which value to assign to the outputted tensor: it's a boolean mask.

For the purpose of this example, I have used random tensors:

>>> a = torch.rand(2, 2, dtype=float)*100
>>> b = torch.rand(2, 2, dtype=float)*0.01
>>> c = torch.rand(2, 2, dtype=float)*10

>>> torch.where(a*(b@c) < 1, -a*b, 0.)
tensor([[ 0.0000,  0.0000],
        [ 0.0000, -0.0183]], dtype=torch.float64)

More generally though, this will work if tensor_a and tensor_b have a shape of (m, n), and tensor_c has a shape of (n, m) because of the operation constraints. In your experiment I'm guessing you only had columns.