How to calculate the median of a masked tensor along an axis?
I have tensor X of floats of dimensions n x m and a tensor Y of booleans of dimensions n x m. I want to calculate values such as the mean, median and max of X, along one of the axes, but only considering the values in X which are true in Y. Something like X[Y].mean(dim=1). This is not possible because X[Y] is always a 1D tensor.
Edit:
For the mean, I was able to do it with:
masked_X = X * Y
masked_X_mean = masked_X.sum(dim=1) / Y.sum(dim=1)

For the max:
masked_X = X
masked_X[Y] = float('-inf')
masked_X_max = masked_X.max(dim=1)

But for the median, I was not able to be as creative. Any suggestions??
e.g.
X = torch.tensor([[1, 1, 1],
                  [2, 2, 4]]).type(torch.float32)
Y = torch.tensor([[0, 1, 0],
                  [1, 0, 1]]).type(torch.bool)

Expected Output
mean = [1., 3.]
median = [1., 2.]
var = [0., 1.]


This is the best I have so far on this:
outs = []
for x, y in zip(X, Y):  # X, Y could be permuted to loop over desired axis
    out = torch.median(torch.masked_select(x, y))
    outs.append(out)
torch.tensor(outs)

Would really appreciate if someone has better solution.
