Fix aten_isclose handling of infinities and equal_nan by Om-singhaI · Pull Request #3026 · microsoft/onnxscript

om singhal (Om-singhaI) · 2026-08-31T01:02:40Z

torch.isclose is not a single tolerance test. In
aten/src/ATen/native/TensorCompare.cpp it is the union of three terms:

close  = (self == other)
close |= isnan(self) & isnan(other)                 # only when equal_nan is set
close |= isfinite(actual_error) & (actual_error <= allowed_error)

The lowering in onnxscript/function_libs/torch_lib/ops/core.py implements the third
term only, and drops the isfinite guard from it:

# FIXME: check equal_nan when self and other are all NaN
# |input - other| <= atol + rtol x |other|
left_part = op.Abs(op.Sub(self, other))
right_part = op.Add(atol, op.Mul(rtol, op.Abs(other)))
result = op.LessOrEqual(left_part, right_part)

So equal_nan never reaches the graph, which the FIXME already noted, and the
infinities come out wrong in both directions, which it did not.

I did not find an issue for this, so there is no number to reference.

What comes out

torch.onnx.export(..., dynamo=True) of torch.isclose(a, b, rtol=1e-05, atol=1e-08)
on main at 3ba2bf7, float32, run under onnxruntime:

         a          b   torch    onnx   equal_nan
       inf        inf    True   False   False
      -inf       -inf    True   False   False
       inf       -inf   False    True   False
         1        inf   False    True   False
        -5       -inf   False    True   False
       nan        nan   False   False   False
       nan        nan    True   False   True
       nan          1   False   False   False
         1   1.000001    True    True   False
         3        3.5   False   False   False

Six of the ten rows disagree. Two mechanisms are at work:

An infinite other makes allowed_error infinite, so |self - other| <= allowed_error
holds for every finite self and for the opposite infinity. Anything measured against
an infinity comes back close.
inf against inf gives actual_error = NaN, and NaN <= inf is false, so two equal
infinities come back not close. The exact equality term is what torch uses to catch
these, and it is missing here.

The attributes have to be bound for any of this to show up. Building a standalone model
from aten_isclose.to_model_proto() leaves rtol and atol as unbound reference
attributes, onnxruntime then evaluates rtol as 0, 0 * inf is NaN, and every row comes
back False, which hides the bug rather than showing it. Every number above came through
torch.onnx.export.

The integer path is unchanged by this. atol and rtol truncate to 0 there, so the
comparison collapses to equality, both before and after. That does differ from torch for a
separate reason, which I describe at the end.

The fix

aten_isclose becomes trace_only=True so that equal_nan and the dtype branch resolve
at export time, and it then follows the three terms directly:

result = op.Equal(self, other)

# |self - other| <= atol + rtol x |other|
actual_error = op.Abs(op.Sub(self, other))
allowed_error = op.Add(
    op.CastLike(atol, other), op.Mul(op.CastLike(rtol, other), op.Abs(other))
)
within_tolerance = op.LessOrEqual(actual_error, allowed_error)

if self.dtype.is_floating_point():
    if equal_nan:
        result = op.Or(result, op.And(op.IsNaN(self), op.IsNaN(other)))
    error = actual_error
    if self.dtype in {ir.DataType.FLOAT16, ir.DataType.BFLOAT16}:
        error = op.Cast(actual_error, to=FLOAT.dtype)
    error_is_finite = op.And(op.Not(op.IsInf(error)), op.Not(op.IsNaN(error)))
    within_tolerance = op.And(error_is_finite, within_tolerance)

return op.Or(result, within_tolerance)

The tolerance band itself is unchanged. The CastLike pair is the same conversion
onnxscript was inserting implicitly for the two float attributes in the scripted version,
written out because a traced function does not insert it. The integer graph therefore
computes exactly what it computed before, with Equal and Or added around it.

Integers are always finite and never NaN, and IsInf and IsNaN do not accept integer
tensors, so the guard and the equal_nan term are skipped for them. That mirrors torch,
which applies the NaN term only when the input is floating point or complex.

ONNX has no IsFinite, so it is spelled the way aten_isfinite spells it a few lines
below in the same file. IsInf accepts only FLOAT and DOUBLE before opset 20, so on
FLOAT16 and BFLOAT16 the error is cast up to FLOAT first. Widening a float is exact, and
the extra node appears only on those two dtypes.

Testing

ops_test_data.py needs no change. The isclose entry carries no skip or xfail, and
sample_inputs_isclose only ever draws finite values out of make_tensor, so the OpInfo
samples cannot reach any of this even though they do sweep equal_nan and rtol.

Added test_isclose_handles_infinities_and_equal_nan and test_isclose_integer_inputs to
tests/function_libs/torch_lib/e2e_ops_tests.py, following the optimize=False pattern
the tests around them use. The first runs the nine pairs from the table above under
float32 and float16 with equal_nan both ways, the second runs int64.

With only the core.py change reverted:

pytest tests/function_libs/torch_lib/e2e_ops_tests.py -k isclose
4 failed, 1 passed, 100 deselected in 6.62s

All four float cases fail with AssertionError: Tensor-likes are not equal!, 5 of 9
elements mismatched with equal_nan off and 6 of 9 with it on. The int64 case passes,
which is the point of it. With the fix:

pytest tests/function_libs/torch_lib/e2e_ops_tests.py -k isclose
5 passed, 100 deselected in 6.06s

pytest tests/function_libs/torch_lib/ops_test.py -k isclose
4 passed, 2 skipped, 1846 deselected, 68 subtests passed in 4.92s

The baseline for the second one on main is
5 passed, 1 skipped, 1846 deselected, 68 subtests passed. The subtest count is
unchanged, so no OpInfo sample moved. The extra skip is the function proto validity
check, which does not apply to traced functions.

Also checked by hand: export plus onnx.checker.check_model(full_check=True) for float32,
float64, float16, bfloat16, int64, int32, int16 and int8. All eight pass the checker, and
all run under onnxruntime except bfloat16, which has no CPU kernel for Equal. bfloat16
had no CPU kernel for Sub before this change either, so it is no worse off.

One thing I left alone

torch computes the tolerance for integer inputs in floating point, because rtol is a
double and rtol * other promotes, so torch.isclose of 999999 against 1000000 is
True with an allowed error of 10. The lowering casts atol and rtol to the input dtype,
where both become 0, and returns False:

torch: [True, True, False]
onnx : [False, True, False]

That holds before and after this change. It is a different root cause from the two missing
terms, and fixing it would move numbers on the integer path, so I kept it out of this
change. Happy to follow up on it separately if you want it.

Environment: Python 3.10, torch 2.9.1, onnx 1.22.0, onnxruntime 1.23.2, macOS arm64.

`torch.isclose` is not a single tolerance test. In `aten/src/ATen/native/TensorCompare.cpp` it is the union of three terms: ``` close = (self == other) close |= isnan(self) & isnan(other) # only when equal_nan is set close |= isfinite(actual_error) & (actual_error <= allowed_error) ``` The lowering in `onnxscript/function_libs/torch_lib/ops/core.py` implements the third term only, and drops the `isfinite` guard from it: ```python # FIXME: check equal_nan when self and other are all NaN # |input - other| <= atol + rtol x |other| left_part = op.Abs(op.Sub(self, other)) right_part = op.Add(atol, op.Mul(rtol, op.Abs(other))) result = op.LessOrEqual(left_part, right_part) ``` So `equal_nan` never reaches the graph, which the FIXME already noted, and the infinities come out wrong in both directions, which it did not. I did not find an issue for this, so there is no number to reference. ## What comes out `torch.onnx.export(..., dynamo=True)` of `torch.isclose(a, b, rtol=1e-05, atol=1e-08)` on main at 3ba2bf7, float32, run under onnxruntime: ``` a b torch onnx equal_nan inf inf True False False -inf -inf True False False inf -inf False True False 1 inf False True False -5 -inf False True False nan nan False False False nan nan True False True nan 1 False False False 1 1.000001 True True False 3 3.5 False False False ``` Six of the ten rows disagree. Two mechanisms are at work: - An infinite `other` makes `allowed_error` infinite, so `|self - other| <= allowed_error` holds for every finite `self` and for the opposite infinity. Anything measured against an infinity comes back close. - `inf` against `inf` gives `actual_error = NaN`, and `NaN <= inf` is false, so two equal infinities come back not close. The exact equality term is what torch uses to catch these, and it is missing here. The attributes have to be bound for any of this to show up. Building a standalone model from `aten_isclose.to_model_proto()` leaves `rtol` and `atol` as unbound reference attributes, onnxruntime then evaluates `rtol` as 0, `0 * inf` is NaN, and every row comes back False, which hides the bug rather than showing it. Every number above came through `torch.onnx.export`. The integer path is not affected. `atol` and `rtol` cast to 0 there, so the comparison collapses to equality, which is what torch returns for the same inputs. ## The fix `aten_isclose` becomes `trace_only=True` so that `equal_nan` and the dtype branch resolve at export time, and it then follows the three terms directly: ```python result = op.Equal(self, other) # |self - other| <= atol + rtol x |other| actual_error = op.Abs(op.Sub(self, other)) allowed_error = op.Add( op.CastLike(atol, other), op.Mul(op.CastLike(rtol, other), op.Abs(other)) ) within_tolerance = op.LessOrEqual(actual_error, allowed_error) if self.dtype.is_floating_point(): if equal_nan: result = op.Or(result, op.And(op.IsNaN(self), op.IsNaN(other))) error = actual_error if self.dtype in {ir.DataType.FLOAT16, ir.DataType.BFLOAT16}: error = op.Cast(actual_error, to=FLOAT.dtype) error_is_finite = op.And(op.Not(op.IsInf(error)), op.Not(op.IsNaN(error))) within_tolerance = op.And(error_is_finite, within_tolerance) return op.Or(result, within_tolerance) ``` The tolerance band itself is unchanged. The `CastLike` pair is the same conversion onnxscript was inserting implicitly for the two float attributes in the scripted version, written out because a traced function does not insert it. The integer graph therefore computes exactly what it computed before, with `Equal` and `Or` added around it. Integers are always finite and never NaN, and `IsInf` and `IsNaN` do not accept integer tensors, so the guard and the `equal_nan` term are skipped for them. That mirrors torch, which applies the NaN term only when the input is floating point or complex. ONNX has no `IsFinite`, so it is spelled the way `aten_isfinite` spells it a few lines below in the same file. `IsInf` accepts only FLOAT and DOUBLE before opset 20, so on FLOAT16 and BFLOAT16 the error is cast up to FLOAT first. Widening a float is exact, and the extra node appears only on those two dtypes. ## Testing `ops_test_data.py` needs no change. The `isclose` entry carries no skip or xfail, and `sample_inputs_isclose` only ever draws finite values out of `make_tensor`, so the OpInfo samples cannot reach any of this even though they do sweep `equal_nan` and `rtol`. Added `test_isclose_handles_infinities_and_equal_nan` and `test_isclose_integer_inputs` to `tests/function_libs/torch_lib/e2e_ops_tests.py`, following the `optimize=False` pattern the tests around them use. The first runs the nine pairs from the table above under float32 and float16 with `equal_nan` both ways, the second runs int64. With only the `core.py` change reverted: ``` pytest tests/function_libs/torch_lib/e2e_ops_tests.py -k isclose 4 failed, 1 passed, 100 deselected in 6.62s ``` All four float cases fail with `AssertionError: Tensor-likes are not equal!`, 5 of 9 elements mismatched with `equal_nan` off and 6 of 9 with it on. The int64 case passes, which is the point of it. With the fix: ``` pytest tests/function_libs/torch_lib/e2e_ops_tests.py -k isclose 5 passed, 100 deselected in 6.06s pytest tests/function_libs/torch_lib/ops_test.py -k isclose 4 passed, 2 skipped, 1846 deselected, 68 subtests passed in 4.92s ``` The baseline for the second one on main is `5 passed, 1 skipped, 1846 deselected, 68 subtests passed`. The subtest count is unchanged, so no OpInfo sample moved. The extra skip is the function proto validity check, which does not apply to traced functions. Also checked by hand: export plus `onnx.checker.check_model(full_check=True)` for float32, float64, float16, bfloat16, int64, int32, int16 and int8. All eight pass the checker, and all run under onnxruntime except bfloat16, which has no CPU kernel for `Equal`. bfloat16 had no CPU kernel for `Sub` before this change either, so it is no worse off. Environment: Python 3.10, torch 2.9.1, onnx 1.22.0, onnxruntime 1.23.2, macOS arm64.

Copilot

Pull request overview

This pull request fixes the ONNX lowering of aten::isclose to match PyTorch semantics for infinities and equal_nan, aligning behavior with the three-term definition used by torch.isclose in PyTorch’s native implementation.

Changes:

Update aten_isclose lowering to include exact equality, optional NaN==NaN handling (equal_nan), and an isfinite(actual_error) guard around the tolerance-band term.
Make aten_isclose trace_only=True so dtype- and equal_nan-dependent branches resolve at export time.
Add E2E tests covering infinities/NaNs for float32/float16 and a regression test ensuring integer behavior is unchanged.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated no comments.

File	Description
`onnxscript/function_libs/torch_lib/ops/core.py`	Reworks `aten_isclose` lowering to match PyTorch’s three-term logic and correct infinity/NaN behavior.
`tests/function_libs/torch_lib/e2e_ops_tests.py`	Adds targeted E2E tests validating the corrected `isclose` behavior for float and integer inputs.

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

codecov · 2026-08-31T17:10:47Z

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 72.67%. Comparing base (3ba2bf7) to head (abb8219).
⚠️ Report is 1 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3026      +/-   ##
==========================================
+ Coverage   72.64%   72.67%   +0.02%     
==========================================
  Files         265      265              
  Lines       32251    32266      +15     
  Branches     3050     3052       +2     
==========================================
+ Hits        23429    23448      +19     
+ Misses       7786     7782       -4     
  Partials     1036     1036

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

Justin Chu (justinchuby) · 2026-08-31T18:48:15Z

+        if self.dtype in {ir.DataType.FLOAT16, ir.DataType.BFLOAT16}:
+            # IsInf takes only FLOAT and DOUBLE before opset 20. Widening is exact.
+            error = op.Cast(actual_error, to=FLOAT.dtype)


This is making the graph quite a bit more complicated for the common path, which I am concerned about.

Is there a way to make the graph, say, only different when equal_nan is true, by making some assumptions on the input values?

You're right, that was too much to put on the common path.

The new shape drops the explicit finiteness check and compares the difference against zero instead. An infinity on either side makes both the error and the allowance infinite, their difference NaN, and every comparison against NaN is false, so the infinities fall out on their own. Equal still carries the case of two equal infinities. IsInf, both Not and both And are gone, and with IsInf goes the Cast that FLOAT16 and BFLOAT16 needed before opset 20, so all four float dtypes now emit the same graph.

On the equal_nan question, not quite. The float path is 14 nodes against main's 10 with equal_nan off, so it still costs four more in the common case: Equal, Or, the extra Sub and the zero constant. Those pay for the infinity behaviour, which is wrong independently of equal_nan, so I do not see a way to make them conditional without leaving inf against inf returning False. equal_nan itself now adds only IsNaN twice plus And and Or, and the integer path is unchanged at 12.

It was 18 and 22 before this commit, so the common path is down by four either way. Elementwise the results are identical to the previous commit: 20736 comparisons, 576 ordered pairs of 24 values across float32, float64 and float16, six rtol and atol pairs, equal_nan both ways, nothing moved.

justinchuby pointed out that the previous shape complicated the common path. It added IsInf, IsNaN, two Not and two And to every exported isclose, plus a Cast on FLOAT16 and BFLOAT16 because IsInf does not take half precision before opset 20. He is right. Comparing the two sides of the tolerance band by subtraction carries the finiteness restriction on its own, so the explicit finiteness check goes away. An infinity on either side makes both the error and the allowance infinite, their difference NaN, and every comparison against NaN false, which leaves the equality term as the only way an infinity is reported close. On finite values the sign of the difference and the direct comparison agree exactly, so the tolerance band itself is unchanged. Exported node counts for float32 with optimize=False: main before after equal_nan False 10 18 14 equal_nan True 10 22 18 FLOAT16 and BFLOAT16 were 19 and 23 because of the extra Cast. They are now 14 and 18, the same as every other float dtype, and the opset 20 caveat is gone with the Cast. The only difference between the two rows above is IsNaN twice plus And and Or, so the graph now changes only when equal_nan is set. The integer path keeps the previous graph at 12 nodes. The subtraction is left off there because it can overflow a narrow integer type where the direct comparison cannot, and the equality term still matters on int8, where a large rtol makes the allowed error wrap. Behavior is unchanged from the previous commit. Sweeping the exported model against torch over 576 ordered pairs drawn from 24 values, for float32, float64 and float16, six rtol and atol pairs, and equal_nan both ways, gives 20736 comparisons. All 20736 match the previous commit element for element, as do 3528 integer comparisons over int64, int32 and int8. A handful of residual disagreements with torch remain on that grid, and how many depends on how the grid samples the edge of the band. They are the same ones as before this commit and are present on main too: they come from the float32 round trip of the two scalar attributes, which puts a value sitting exactly on the boundary just outside it, not from the graph shape. pytest tests/function_libs/torch_lib/e2e_ops_tests.py -k isclose 5 passed, 100 deselected pytest tests/function_libs/torch_lib/ops_test.py -k isclose 4 passed, 2 skipped, 1846 deselected onnx.checker.check_model(full_check=True) still passes for float32, float64, float16, bfloat16, int64, int32, int16 and int8, with equal_nan both ways.

Copilot

Pull request overview

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

…triction

Copilot

Pull request overview

Copilot reviewed 2 out of 2 changed files in this pull request and generated no new comments.

github-project-automation Bot added this to ONNX Script Review Board Aug 31, 2026

github-project-automation Bot moved this to Todo in ONNX Script Review Board Aug 31, 2026

Justin Chu (justinchuby) requested a lite review from Copilot August 31, 2026 17:01

Justin Chu (justinchuby) added the module: torchlib Related to the torch/aten function lib in development label Aug 31, 2026

Copilot started reviewing on behalf of Justin Chu (justinchuby) August 31, 2026 17:07 View session

Copilot AI reviewed Aug 31, 2026

View reviewed changes

Justin Chu (justinchuby) reviewed Aug 31, 2026

View reviewed changes

Justin Chu (justinchuby) self-assigned this Aug 31, 2026

Justin Chu (justinchuby) requested a lite review from Copilot September 1, 2026 01:47

Copilot started reviewing on behalf of Justin Chu (justinchuby) September 1, 2026 01:47 View session

Copilot AI reviewed Sep 1, 2026

View reviewed changes

Comment thread onnxscript/function_libs/torch_lib/ops/core.py Outdated

Correct the comment on why the subtraction carries the finiteness res…

abb8219

…triction

Justin Chu (justinchuby) requested a lite review from Copilot September 1, 2026 18:37

Copilot started reviewing on behalf of Justin Chu (justinchuby) September 1, 2026 18:37 View session

Copilot AI reviewed Sep 1, 2026

View reviewed changes

Navigation Menu

Uh oh!

Fix aten_isclose handling of infinities and equal_nan - #3026

Fix aten_isclose handling of infinities and equal_nan#3026
om singhal (Om-singhaI) wants to merge 3 commits into
microsoft:mainfrom
Om-singhaI:fix/isclose-inf-and-equal-nan

om singhal (Om-singhaI) commented Aug 31, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

codecov Bot commented Aug 31, 2026 •

edited

Loading

Uh oh!

Justin Chu (justinchuby) Aug 31, 2026

Uh oh!

Justin Chu (justinchuby) Aug 31, 2026 •

edited

Loading

Uh oh!

om singhal (Om-singhaI) Aug 31, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Copilot AI left a comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

om singhal (Om-singhaI) commented Aug 31, 2026

What comes out

The fix

Testing

One thing I left alone

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

codecov Bot commented Aug 31, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Justin Chu (justinchuby) Aug 31, 2026

Choose a reason for hiding this comment

Uh oh!

Justin Chu (justinchuby) Aug 31, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

om singhal (Om-singhaI) Aug 31, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

codecov Bot commented Aug 31, 2026 •

edited

Loading

Justin Chu (justinchuby) Aug 31, 2026 •

edited

Loading