[Shawn (ShawnNotFound)]

torch.fmax, torch.fmin and torch.hypot currently fail to export:

class Model(torch.nn.Module):
    def forward(self, a, b):
        return torch.fmax(a, b)

torch.onnx.export(Model(), (torch.randn(3), torch.randn(3)), dynamo=True)
# DispatchError: No ONNX function found for <OpOverload(op='prims.fmax', overload='default')>

This implements them and registers them for both the aten and prims overloads, like aten_remainder does.

• fmax / fmin: Max and Min return NaN when either input is NaN, but fmax and fmin return the other value and only give NaN when both inputs are NaN. For floating point inputs the NaN positions are filled from the other input. Bool inputs use Or / And as in aten_maximum / aten_minimum.
• hypot: computing sqrt(x*x + y*y) directly overflows or underflows well before the result does, e.g. hypot(300, 400) in float16. It's computed as big * sqrt(1 + (small / big)**2) instead, with zeros and infinities handled separately so that it matches torch.hypot, including hypot(inf, nan) == inf. I used Abs(x) == inf rather than IsInf, since IsInf only accepts float16 from opset 20.

Tests: OpInfo entries for the three ops, and e2e tests for NaN inputs, integer inputs and hypot inputs whose squares overflow or underflow. Against PyTorch on random inputs over a wide range, hypot is within 2 ulp for float32 and float64 and 1 ulp for float16.

[Justin Chu (justinchuby)]

Thanks!

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 73.39%. Comparing base (5b210da) to head (6e44eea).
⚠️ Report is 1 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3169      +/-   ##
==========================================
+ Coverage   73.34%   73.39%   +0.04%     
==========================================
  Files         269      269              
  Lines       33564    33597      +33     
  Branches     3210     3215       +5     
==========================================
+ Hits        24618    24657      +39     
+ Misses       7876     7870       -6     
  Partials     1070     1070              

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

[Copilot]

🟡 Changes recommended

The float16 hypot calculation can overflow to infinity when the expected result is finite, requiring float32 intermediates and a boundary-case regression test.

1 open issue

• Promote float16 hypot intermediates to prevent overflow · New

What changed in this PR

Adds TorchLib export support for torch.fmax, torch.fmin, and torch.hypot, registered for both ATen and prims overloads.

Changes:

• Implements NaN-aware fmax and fmin, including boolean handling.
• Implements scaled hypot computation with zero and infinity handling.
• Adds OpInfo coverage and end-to-end tests for special values, integer inputs, and numerical range.

File	Description
tests/​function_libs/​torch_lib/​ops_test_data.py	Registers OpInfo tests for all three operators.
tests/​function_libs/​torch_lib/​e2e_ops_tests.py	Adds export tests for edge cases and numerical range.
onnxscript/​function_libs/​torch_lib/​ops/​core.py	Implements and registers the three operators.

_{Review effort: Balanced}