Tensor Mean of values greater than a threshold

Question

I have a tensor of shape = [a,b,c]

The tensor mean along dim=1 will give me an output of shape = [a,c]

My goal is to compute the mean of values along dim=1 greater than a threshold.
How is this possible in TensorFlow?

I am aware of tf.where() and tf.boolean_mask() functions but not sure exactly how to use them for this task.

A Kareem · Answer

First, If you calculate the mean along dim=1 the output shape should be [a, c].

If you want to mask the mean that's less then a threshold and set it to zero you can do

# generate data
torch.manual_seed(42)
a, b, c = 2, 3, 4
t = torch.normal(torch.ones(a, b, c))
print(t)

tensor([[[ 2.9269,  2.4873,  1.9007, -1.1055],
         [ 1.6784, -0.2345,  0.9569, -0.6047],
         [ 1.3559,  0.3134,  0.5066,  1.2415]],

[[-0.1109,  1.0915, -1.3169,  0.7832],
         [ 0.6903,  0.6043,  1.8034,  0.3784],
         [ 0.4080,  0.9369,  0.1714,  1.3309]]])

to get the mean and mask to 0

mean = t.mean(dim=1)
print(mean)

tensor([[ 1.9871,  0.8554,  1.1214, -0.1562],
        [ 0.3291,  0.8776,  0.2193,  0.8308]])

threshold = 1
mean[mean<threshold] = 0
print(mean)

tensor([[1.9871, 0.0000, 1.1214, 0.0000],
        [0.0000, 0.0000, 0.0000, 0.0000]])

If you instead just want a list of means you can do this:

print(mean[mean<threshold])

tensor([1.9871, 1.1214])

One Answer