Torch.divide only where denominator is non-zero

In numpy I can do the following to avoid division by zero:

a = np.random.randint(0, 10, 100)
b = np.random.randint(0, 10, 100)
c = np.zeros_like(a, dtype=np.float32) # It can be anything other than zero
c = np.divide(a, b, out=c, where=(b!=0))

In torch.divide there is no where argument for masking. Only way seems to be replacing inf with desired value after the division takes place. Is there any better way to do this in pytorch?

1 Like

Would selecting the desired values in a and b before applying the division work?
E.g. you could create a mask first and use it to index both tensors where the condition would be b!=0.

2 Likes

Yes, indeed your(@ptrblck) solution works pretty well. For anyone who’s looking for solution using torch see the snippet below:

import torch

# numerator: tensor([2., 2., 0., 5., 7., 3., 4., 3., 6., 5.])
a = torch.randint(0, 10, (10,), dtype=torch.float32)

# denominator: tensor([3., 3., 0., 4., 5., 4., 7., 8., 0., 4.])
b = torch.randint(0, 10, (10,), dtype=torch.float32)

# initialize output tensor with desired value
c = torch.full_like(a, fill_value=float('nan'))

# zero mask
mask = (b != 0)

# finally perform division
c[mask] = a[mask] / b[mask]

# output: tensor([0.6667, 0.6667,    nan, 1.2500, 1.4000, 0.7500, 0.5714, 0.3750,    nan, 1.2500])

3 Likes