-
Notifications
You must be signed in to change notification settings - Fork 2.8k
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
base: main
Are you sure you want to change the base?
Conversation
result = lax.div(x1, x2) | ||
result = lax.select(_broadcast_to(x2 == 0, result.shape), | ||
lax.full_like(result, fill_value=0), | ||
result) | ||
return result |
There was a problem hiding this comment.
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:
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)) |
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 |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Similarly here:
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) |
div = lax.select(_broadcast_to(x2 == 0, div.shape), | ||
lax.full_like(div, fill_value=np.inf) * lax.sign(x1), | ||
div) | ||
|
There was a problem hiding this comment.
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).
Fixes #11990
Current behaviour:
With thi PR: