How to calculate torch.mean with default value for empty slice?

user_913098 0
Viewed 123

What's the correct way to calculate torch.mean such that in case of an empty slice, some default value would be returned?

import torch

a = torch.arange(5, dtype=torch.float)
mask = torch.ones(5, dtype=torch.bool)
mask[2] = 0
mask_empty = torch.zeros(5, dtype=torch.bool)

a_masked = a[mask]
a_empty = a[mask_empty]
m = torch.mean(a_masked)
m_empty = torch.mean(a_empty)

print(m)
print(m_empty)

current output:

tensor(2.)
tensor(nan)

required output:

tensor(2.)
tensor(0.)

I know I can do

m_empty = torch.zeros_like(m_empty) if torch.isnan(m_empty) else m_empty

but this seems like the wrong way to go, as it forces more cpu-gpu communication that I would like to avoid if possible.
Also, this makes the code less clean than if something like m_empty = torch.mean(a_empty, default=0) were available.

Is there a clean, Pythonic way of achieving default values for empty slices of torch.mean [or other standard functions]?

pytorch python
1 Answers

This may not be the best way and is not as clean as what you are looking for however, one way to achieve this is to check the length of the Tensor prior to calculating the mean,

import torch

a = torch.arange(5, dtype=torch.float)
mask = torch.ones(5, dtype=torch.bool)
mask[2] = 0
mask_empty = torch.zeros(5, dtype=torch.bool)
a_masked = a[mask]
a_empty = a[mask_empty]
m = torch.mean(a_masked) if len(a_masked) > 0 else torch.tensor(0.)
m_empty = torch.mean(a_empty) if len(a_empty) > 0 else torch.tensor(0.)

print(m)
print(m_empty)

Output,

tensor(2.)
tensor(0.)
user_84587120
Related