Conditionally apply tensor operations in PyTorch
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) &lt; 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.

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 &quot;simple matrices&quot;, 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:
&gt;&gt;&gt; a = torch.rand(2, 2, dtype=float)*100
&gt;&gt;&gt; b = torch.rand(2, 2, dtype=float)*0.01
&gt;&gt;&gt; c = torch.rand(2, 2, dtype=float)*10

&gt;&gt;&gt; torch.where(a*(b@c) &lt; 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.
