atan2: propagate NaN instead of returning 0/pi at x == +/-0#35
Merged
gokulkrishna98 merged 2 commits intoJul 8, 2026
Merged
Conversation
atan2's x=0 branch classified purely via comparisons, which are all False for NaN and misclassified atan2(NaN, +0) as 0 and atan2(NaN, -0) as pi instead of NaN. Add an explicit NaN check (y != y or x != x) that takes priority over every other branch, plus a regression test and updated IR FileCheck patterns.
gokulkrishna98
marked this pull request as ready for review
July 8, 2026 21:01
jakesabathia2
approved these changes
Jul 8, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
The
atan2converter's x=0 branch classifies inputs purely via comparisons (y > 0,0 > y, etc.), which are allFalsewhen NaN is involved. This meantatan2(NaN, +0.0)incorrectly returned0andatan2(NaN, -0.0)incorrectly returnedpi, instead ofNaNas required by IEEE-754 and matched bytorch.atan2.This adds an explicit NaN check (
y != y or x != x) that takes priority over every other branch in the decomposition, so NaN now always propagates through correctly.Testing
test_nan_propagationtoTestAtan2intests/ops/test_ops.py, covering NaN iny,x, and both, including atx = +/-0.TestAtan2IRFileCheck patterns intests/ops/test_ops_ir.pyto match the new IR shape (extra NaN constant + final NaN-selectwhere).tests/ops/suite passes locally.