Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Update jnp.floor_divide to make it consistent with np.floor_divide for division by zero #25032

Open
wants to merge 1 commit into
base: main
Choose a base branch
from

Conversation

rajasekharporeddy
Copy link
Contributor

@rajasekharporeddy rajasekharporeddy commented Nov 21, 2024

Fixes #11990

Current behaviour:

>>> jnp.floor_divide(2, 0)
Array(-2, dtype=int32, weak_type=True)
>>> np.floor_divide(2, 0)
<stdin>:1: RuntimeWarning: divide by zero encountered in floor_divide
0

With thi PR:

>>> jnp.floor_divide(2, 0)
Array(0, dtype=int32, weak_type=True)

@dfm dfm requested a review from jakevdp November 21, 2024 15:42
Comment on lines +2498 to +2502
result = lax.div(x1, x2)
result = lax.select(_broadcast_to(x2 == 0, result.shape),
lax.full_like(result, fill_value=0),
result)
return result
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Using _where would make this more succinct. I think this should work:

Suggested change
result = lax.div(x1, x2)
result = lax.select(_broadcast_to(x2 == 0, result.shape),
lax.full_like(result, fill_value=0),
result)
return result
return _where(x2 == 0, 0, lax.div(x1, x2))

Comment on lines +2507 to +2511
result = _where(select, quotient - 1, quotient)
result = lax.select(_broadcast_to(x2 == 0, result.shape),
lax.full_like(result, fill_value=0),
result)
return result
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Similarly here:

Suggested change
result = _where(select, quotient - 1, quotient)
result = lax.select(_broadcast_to(x2 == 0, result.shape),
lax.full_like(result, fill_value=0),
result)
return result
return _where(x2 == 0, 0, _where(select, quotient - 1, quotient)

Comment on lines +2570 to +2573
div = lax.select(_broadcast_to(x2 == 0, div.shape),
lax.full_like(div, fill_value=np.inf) * lax.sign(x1),
div)

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We need to double check this logic for when x2 is -0.0 (the sign of the returned infinity may be wrong), or when x1 and x2 are both zero (when the result should be NaN).

@jakevdp jakevdp self-assigned this Nov 21, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

BUG: floor_divide returns strange values for integer division by zero
2 participants