I would like to get the mean and standard deviation from a tensor with shape (H,W) along the dimension 1 (so the output would be a pair of tensors with shape (H) or (H,1)). My problem is: I need to filter each row and only use the values selected by a given mask and the mask selects a different number of values from each row.

My actual approach is to generate another tensor with NaNs where I don’t care about the value and use torch.nanmean to get the mean ignoring those values but I don’t find an analogue function to get the std.

This should be feasible (a mean from the values greater than 0) but there is no nanstd function to get the standard deviation ignoring nans. It’s important to note that the first value comes from the mean over 3 values while the second comes from the mean over 2 values (that is the reason I cannot simply filter to get a tensor only with the desired values, the rows length would not match).

Sometimes, you might want to set keepdim=True, here is modified version:

def nanstd(o, dim, keepdim=False):
result = torch.sqrt(
torch.nanmean(
torch.pow( torch.abs(o-torch.nanmean(o,dim=dim).unsqueeze(dim)),2),
dim=dim
)
)
if keepdim:
result = result.unsqueeze(dim)
return result