Create a mask that is larger than the n-th quantile of each 2D tensor in a batch

1

I have a torch.Tensor of shape (2, 2, 2) (can be bigger), where the values are normalized within range [0, 1].

Now I am given a positive integer K, which tells me that I need to create a mask where for each 2D tensor inside the batch, values are 1 if it is larger than 1/k of all the values, and 0 elsewhere. The return mask also has shape (2, 2, 2).

For example, if I have a batch like this:

tensor([[[1., 3.],
         [2., 4.]],
        [[5., 7.],
         [9., 8.]]])

and let K=2, it means that I must mask the values where they are greater than 50% of all the values inside each 2D tensor.

In the example, the 0.5 quantile is 2.5 and 7.5, so this is the desired output:

tensor([[[0, 1],
         [0, 1]],
        [[0, 0],
         [1, 1]]])

I tried:

a = torch.tensor([[[0, 1],
                   [0, 1]],
                  [[0, 0],
                   [1, 1]]])
quantile = torch.tensor([torch.quantile(x, 1/K) for x in a])
torch.where(a > val, 1, 0)

But this is the result:

tensor([[[0, 0],
         [0, 0]],
        [[1, 0],
         [1, 1]]])
2

You could use broadcasting to get the desired result:

x = torch.tensor([[[1., 3.],
                   [2., 4.]],
                  [[5., 7.],
                   [9., 8.]]])

quantile = torch.tensor([torch.quantile(a, 1/2) for a in x])

res = torch.where(x > quantile[:, None, None], 1, 0)
print(res)
> tensor([[[0, 1],
           [0, 1]],

          [[0, 0],
           [1, 1]]])