Implemented missing torch.nan* operators

Open in Colab

import torch
from maskedtensor import masked_tensor

Issue 21987

This issue was closed by inclusion into Issue 61474 - Implement missing torch.nan* operators. This proposes an alternative, which is to use masked tensors instead of introducing additional operators. Since nanmean has already landed, we can use it as a comparison point.

y = torch.arange(32).float()
x = y * y.fmod(4)
x = x.masked_fill(x == 0, float('nan'))
print(x)
print(torch.nanmean(x))
print(torch.mean(masked_tensor(x, ~torch.isnan(x))))

tensor([nan,  1.,  4.,  9., nan,  5., 12., 21., nan,  9., 20., 33., nan, 13.,
        28., 45., nan, 17., 36., 57., nan, 21., 44., 69., nan, 25., 52., 81.,
        nan, 29., 60., 93.])
tensor(32.6667)
masked_tensor( 32.6667, True)

MaskedTensor can further support reduction when fully masked out, as would be the case when a given Tensor is completetely nan. nanmean on the other hand returns nan when the input is entirely nan.

x = torch.empty(32)
x.fill_(float('nan'))
print(x)
print(torch.nanmean(x))
print(torch.mean(masked_tensor(x, ~torch.isnan(x))))

tensor([nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
        nan, nan, nan, nan, nan, nan, nan, nan])
tensor(nan)
masked_tensor(--, False)

Further some users already want to use nan reductions to encode masked semantics.