Apply operation on tensor with masked value

Let’s say you have a tensor like this:
mytensor = torch.tensor([[[-11,0,101,0],[0,5,50,0],[-1,-2,0,1.]]])
And you define your mask as being 0: mask = mytensor.eq(0.)
Now you want to apply in dim 1 several metrics such has max, min, mean (and potentially custom one).
Doing something like this:

mytensor.mean(1)
tensor([[-4.0000,  1.0000, 50.3333,  0.3333]])

Is not going to work for example for the first value it should be (-11-1)/2=-6
You can hack getting the mean like this: mytensor.sum(1).div((~mask).sum(1).float()) but this is not going to work for min and max or any custom operation.
nn.CrossEntropyLoss has a parameter ignore_index to resolve that but I can’t see any way to do it for generic operation. Is there a clean non hacky way to do it?

1 Like

Any updates on this?

custom metrics may be complex.
when min, you can change 0 element in mytensor to inf;
when max, you can change 0 element in mytensor to -inf;
There may be some more easy way to complete your task.

1 Like

Yes, using Inf’s is theoretically possible, but I am getting NaN’s in gradients afterwards… That’s why I wonder if that’s a safe solution…

sorry, I do not consider the gradients.

>>> import torch
>>> mytensor = torch.tensor([[-11, 0, 101, 0], [0, 5, 50, 0], [-1, -2, 0, 1]])
>>> mask = mytensor.eq(0.)
>>> mytensor_sum = mytensor.sum(1)
>>> mytensor_sum = mytensor_sum.unsqueeze(1)
>>> other = mytensor_sum.repeat(1, 4)
>>> other
tensor([[ 90,  90,  90,  90],
        [ 55,  55,  55,  55],
        [ -2,  -2,  -2,  -2]])
>>> final = torch.where(mask, other, mytensor)
>>> final
tensor([[ -11,   90,  101,   90],
        [  55,    5,   50,   55],
        [  -1,   -2,   -2,    1]])
>>> fianl.max(1)
(tensor([ 101,   55,    1]), tensor([ 2,  0,  3]))
>>> final.min(1)
(tensor([-11,   5,  -2]), tensor([ 0,  1,  1]))

I believe that others will have a simpler approach.