Torch.divide only where denominator is non-zero

1

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
2

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.

3 Likes
3

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])
4 Likes
4

Another solution is to use where before you divide (this avoids initialising c and then conditionally modifying its elements in-place - I think this looks more elegant, but I haven’t measured its efficiency):

good_inds = (b != 0)
fill_value = 42 # or whatever
c = (
    torch.where(good_inds, a, fill_value) /
    torch.where(good_inds, b, 1)
)