From 9417a890dfa93b0669b76d669a71c78feba04bf1 Mon Sep 17 00:00:00 2001 From: ShawnNotFound Date: Mon, 2 Nov 2026 22:34:38 +0000 Subject: [PATCH 1/3] Implement fmax and fmin Max and Min return NaN when either input is NaN, so for floating point inputs the other value is selected where one input is NaN. Both ops are also registered for the prims overloads, which is what torch.export produces for torch.fmax and torch.fmin. --- .../function_libs/torch_lib/ops/core.py | 30 +++++++++++++--- .../function_libs/torch_lib/e2e_ops_tests.py | 36 +++++++++++++++++++ .../function_libs/torch_lib/ops_test_data.py | 2 ++ 3 files changed, 64 insertions(+), 4 deletions(-) diff --git a/onnxscript/function_libs/torch_lib/ops/core.py b/onnxscript/function_libs/torch_lib/ops/core.py index 86652cafb7..f395c40560 100644 --- a/onnxscript/function_libs/torch_lib/ops/core.py +++ b/onnxscript/function_libs/torch_lib/ops/core.py @@ -4238,16 +4238,38 @@ def operator_floordiv(self: INT64, other: INT64) -> INT64: return op.Div(self, other) -def aten_fmax(self: TensorType, other: TensorType) -> TensorType: +@torch_op(("aten::fmax", "prims::fmax"), trace_only=True) +def aten_fmax(self: TTensor, other: TTensor) -> TTensor: """fmax(Tensor self, Tensor other) -> Tensor""" - raise NotImplementedError() + if self.dtype == ir.DataType.BOOL: + return op.Or(self, other) + + result = op.Max(self, other) + if not self.dtype.is_floating_point(): + return result + + # Max returns NaN if either input is NaN. fmax returns the other input instead, + # and NaN only when both are NaN. + result = op.Where(op.IsNaN(self), other, result) + return op.Where(op.IsNaN(other), self, result) -def aten_fmin(self: TensorType, other: TensorType) -> TensorType: +@torch_op(("aten::fmin", "prims::fmin"), trace_only=True) +def aten_fmin(self: TTensor, other: TTensor) -> TTensor: """fmin(Tensor self, Tensor other) -> Tensor""" - raise NotImplementedError() + if self.dtype == ir.DataType.BOOL: + return op.And(self, other) + + result = op.Min(self, other) + if not self.dtype.is_floating_point(): + return result + + # Min returns NaN if either input is NaN. fmin returns the other input instead, + # and NaN only when both are NaN. + result = op.Where(op.IsNaN(self), other, result) + return op.Where(op.IsNaN(other), self, result) @torch_op(("aten::fmod.Tensor", "aten::fmod.Scalar"), trace_only=True) diff --git a/tests/function_libs/torch_lib/e2e_ops_tests.py b/tests/function_libs/torch_lib/e2e_ops_tests.py index 5c88e927f6..f534975d68 100644 --- a/tests/function_libs/torch_lib/e2e_ops_tests.py +++ b/tests/function_libs/torch_lib/e2e_ops_tests.py @@ -2013,6 +2013,42 @@ def forward(self, a, b): onnx_program = torch.onnx.export(IsCloseModel(), (a, b), dynamo=True, optimize=False) _testing.assert_onnx_program(onnx_program) + @parameterized.parameterized.expand( + [ + ("fmax_float32", "fmax", torch.float32), + ("fmax_float16", "fmax", torch.float16), + ("fmin_float32", "fmin", torch.float32), + ("fmin_float16", "fmin", torch.float16), + ] + ) + def test_fmax_fmin_ignore_nan(self, _: str, func: str, dtype: torch.dtype): + # fmax and fmin return the other input when one of them is NaN, and NaN + # only when both are. + func = getattr(torch, func) + + class Model(torch.nn.Module): + def forward(self, a, b): + return func(a, b) + + nan = math.nan + inf = math.inf + a = torch.tensor([nan, 1.0, nan, -inf, inf, 2.0, -3.0], dtype=dtype) + b = torch.tensor([1.0, nan, nan, nan, 0.5, 5.0, -4.0], dtype=dtype) + + onnx_program = torch.onnx.export(Model(), (a, b), dynamo=True, optimize=False) + _testing.assert_onnx_program(onnx_program) + + def test_fmax_fmin_integer_inputs(self): + class Model(torch.nn.Module): + def forward(self, a, b): + return torch.fmax(a, b), torch.fmin(a, b) + + a = torch.tensor([1, -2, 3, 7], dtype=torch.int64) + b = torch.tensor([4, -5, 3, 0], dtype=torch.int64) + + onnx_program = torch.onnx.export(Model(), (a, b), dynamo=True, optimize=False) + _testing.assert_onnx_program(onnx_program) + @parameterized.parameterized.expand( [ ("amax", "amax", False), diff --git a/tests/function_libs/torch_lib/ops_test_data.py b/tests/function_libs/torch_lib/ops_test_data.py index 15a0cec84c..d4de4cbdb9 100644 --- a/tests/function_libs/torch_lib/ops_test_data.py +++ b/tests/function_libs/torch_lib/ops_test_data.py @@ -732,6 +732,8 @@ def _where_input_wrangler( TorchLibOpInfo("flatten", core_ops.aten_flatten), TorchLibOpInfo("floor", core_ops.aten_floor), TorchLibOpInfo("ops.aten.floor_divide", core_ops.aten_floor_divide), + TorchLibOpInfo("fmax", core_ops.aten_fmax), + TorchLibOpInfo("fmin", core_ops.aten_fmin), TorchLibOpInfo("fmod", core_ops.aten_fmod), TorchLibOpInfo("frac", core_ops.aten_frac), TorchLibOpInfo("full", core_ops.aten_full), From 7f39d49719d1ab3022e212b901027faee46f9287 Mon Sep 17 00:00:00 2001 From: ShawnNotFound Date: Mon, 2 Nov 2026 22:36:55 +0000 Subject: [PATCH 2/3] Implement hypot Squaring the inputs overflows or underflows well before the result does, e.g. hypot(300, 400) in float16, so the larger magnitude is factored out first. Zero and infinite inputs are handled separately so that the results match torch.hypot, including hypot(inf, nan) == inf. --- .../function_libs/torch_lib/ops/core.py | 24 +++++++++++++++-- .../function_libs/torch_lib/e2e_ops_tests.py | 27 +++++++++++++++++++ .../function_libs/torch_lib/ops_test_data.py | 1 + 3 files changed, 50 insertions(+), 2 deletions(-) diff --git a/onnxscript/function_libs/torch_lib/ops/core.py b/onnxscript/function_libs/torch_lib/ops/core.py index f395c40560..3412f5c64c 100644 --- a/onnxscript/function_libs/torch_lib/ops/core.py +++ b/onnxscript/function_libs/torch_lib/ops/core.py @@ -5004,10 +5004,30 @@ def reshape_to_atleast_2d(tensor): return result -def aten_hypot(self: TensorType, other: TensorType) -> TensorType: +@torch_op(("aten::hypot", "prims::hypot"), trace_only=True) +def aten_hypot(self: TFloat, other: TFloat) -> TFloat: """hypot(Tensor self, Tensor other) -> Tensor""" - raise NotImplementedError() + zero = common_ops.constant(0.0, dtype=self.dtype) + one = common_ops.constant(1.0, dtype=self.dtype) + inf = common_ops.constant(math.inf, dtype=self.dtype) + + # Squaring the inputs directly overflows or underflows long before the result + # does, e.g. hypot(300, 400) in float16. Scale by the larger magnitude instead: + # hypot(x, y) = big * sqrt(1 + (small / big) ** 2) + abs_self = op.Abs(self) + abs_other = op.Abs(other) + big = op.Max(abs_self, abs_other) + small = op.Min(abs_self, abs_other) + ratio = op.Div(small, big) + result = op.Mul(big, op.Sqrt(op.Add(one, op.Mul(ratio, ratio)))) + + # 0 / 0 and inf / inf give NaN above. hypot is 0 when both inputs are 0, and + # inf when either input is inf, even if the other one is NaN. + result = op.Where(op.Equal(big, zero), zero, result) + # IsInf only accepts float16 from opset 20, so compare the magnitudes instead. + either_inf = op.Or(op.Equal(abs_self, inf), op.Equal(abs_other, inf)) + return op.Where(either_inf, inf, result) def aten_i0(self: TensorType) -> TensorType: diff --git a/tests/function_libs/torch_lib/e2e_ops_tests.py b/tests/function_libs/torch_lib/e2e_ops_tests.py index f534975d68..130d624af8 100644 --- a/tests/function_libs/torch_lib/e2e_ops_tests.py +++ b/tests/function_libs/torch_lib/e2e_ops_tests.py @@ -2049,6 +2049,33 @@ def forward(self, a, b): onnx_program = torch.onnx.export(Model(), (a, b), dynamo=True, optimize=False) _testing.assert_onnx_program(onnx_program) + @parameterized.parameterized.expand( + [ + ("float32", torch.float32), + ("float16", torch.float16), + ] + ) + def test_hypot_large_small_and_special_values(self, _: str, dtype: torch.dtype): + # Squaring these inputs overflows or underflows, but the results are in range. + # hypot is inf when either input is inf, even if the other one is NaN. + class Model(torch.nn.Module): + def forward(self, a, b): + return torch.hypot(a, b) + + nan = math.nan + inf = math.inf + big = torch.finfo(dtype).max / 4 + tiny = torch.finfo(dtype).tiny * 4 + a = torch.tensor( + [300.0, -big, tiny, 0.0, -0.0, inf, nan, -inf, nan, 3.0, 1.0], dtype=dtype + ) + b = torch.tensor( + [400.0, big, -tiny, 0.0, 5.0, nan, -inf, 2.0, 1.0, -4.0, 1e-3], dtype=dtype + ) + + onnx_program = torch.onnx.export(Model(), (a, b), dynamo=True, optimize=False) + _testing.assert_onnx_program(onnx_program) + @parameterized.parameterized.expand( [ ("amax", "amax", False), diff --git a/tests/function_libs/torch_lib/ops_test_data.py b/tests/function_libs/torch_lib/ops_test_data.py index d4de4cbdb9..fa8890b4ef 100644 --- a/tests/function_libs/torch_lib/ops_test_data.py +++ b/tests/function_libs/torch_lib/ops_test_data.py @@ -766,6 +766,7 @@ def _where_input_wrangler( dtypes=(torch.float16,), reason="we skip float16 because of https://github.com/pytorch/pytorch/issues/174668", ), + TorchLibOpInfo("hypot", core_ops.aten_hypot), # TorchLibOpInfo("is_same_size", core_ops.aten_is_same_size), # no test case in OPS_DB # TorchLibOpInfo("is_nonzero", core_ops.aten_is_nonzero), # no test case in OPS_DB TorchLibOpInfo("ops.aten.index.Tensor", core_ops.aten_index), From d6fa88f787b63be2a6734598ed34b6f226c1e788 Mon Sep 17 00:00:00 2001 From: ShawnNotFound Date: Fri, 6 Nov 2026 18:32:38 +0000 Subject: [PATCH 3/3] [torchlib] Compute float16 and bfloat16 hypot in float32 Doing the scaling in float32 keeps the intermediate steps in full precision for results near the top of the half precision range. It also makes bfloat16 work on onnxruntime, which has no bfloat16 kernel for Abs. --- .../function_libs/torch_lib/ops/core.py | 24 ++++++++++++++----- .../function_libs/torch_lib/e2e_ops_tests.py | 21 ++++++++++++++++ 2 files changed, 39 insertions(+), 6 deletions(-) diff --git a/onnxscript/function_libs/torch_lib/ops/core.py b/onnxscript/function_libs/torch_lib/ops/core.py index 3412f5c64c..46ad5158e6 100644 --- a/onnxscript/function_libs/torch_lib/ops/core.py +++ b/onnxscript/function_libs/torch_lib/ops/core.py @@ -5008,12 +5008,25 @@ def reshape_to_atleast_2d(tensor): def aten_hypot(self: TFloat, other: TFloat) -> TFloat: """hypot(Tensor self, Tensor other) -> Tensor""" - zero = common_ops.constant(0.0, dtype=self.dtype) - one = common_ops.constant(1.0, dtype=self.dtype) - inf = common_ops.constant(math.inf, dtype=self.dtype) + dtype = self.dtype + if dtype in {ir.DataType.FLOAT16, ir.DataType.BFLOAT16}: + # Compute in float32 so the intermediate steps keep full precision near the + # top of the half precision range. onnxruntime also lacks bfloat16 kernels + # for some of the ops below, e.g. Abs. + self = op.Cast(self, to=FLOAT.dtype) + other = op.Cast(other, to=FLOAT.dtype) + return op.Cast(_hypot(self, other, ir.DataType.FLOAT), to=dtype) + + return _hypot(self, other, dtype) + + +def _hypot(self: TFloat, other: TFloat, dtype: ir.DataType) -> TFloat: + zero = common_ops.constant(0.0, dtype=dtype) + one = common_ops.constant(1.0, dtype=dtype) + inf = common_ops.constant(math.inf, dtype=dtype) # Squaring the inputs directly overflows or underflows long before the result - # does, e.g. hypot(300, 400) in float16. Scale by the larger magnitude instead: + # does, e.g. hypot(1e30, 1e30) in float32. Scale by the larger magnitude instead: # hypot(x, y) = big * sqrt(1 + (small / big) ** 2) abs_self = op.Abs(self) abs_other = op.Abs(other) @@ -5025,8 +5038,7 @@ def aten_hypot(self: TFloat, other: TFloat) -> TFloat: # 0 / 0 and inf / inf give NaN above. hypot is 0 when both inputs are 0, and # inf when either input is inf, even if the other one is NaN. result = op.Where(op.Equal(big, zero), zero, result) - # IsInf only accepts float16 from opset 20, so compare the magnitudes instead. - either_inf = op.Or(op.Equal(abs_self, inf), op.Equal(abs_other, inf)) + either_inf = op.Or(op.IsInf(self), op.IsInf(other)) return op.Where(either_inf, inf, result) diff --git a/tests/function_libs/torch_lib/e2e_ops_tests.py b/tests/function_libs/torch_lib/e2e_ops_tests.py index 130d624af8..91069573fb 100644 --- a/tests/function_libs/torch_lib/e2e_ops_tests.py +++ b/tests/function_libs/torch_lib/e2e_ops_tests.py @@ -2076,6 +2076,27 @@ def forward(self, a, b): onnx_program = torch.onnx.export(Model(), (a, b), dynamo=True, optimize=False) _testing.assert_onnx_program(onnx_program) + @parameterized.parameterized.expand( + [ + ("float16", torch.float16), + ("bfloat16", torch.bfloat16), + ] + ) + def test_hypot_half_results_near_max(self, _: str, dtype: torch.dtype): + # These results are finite but close to the largest value of the type, so + # rounding in a half precision computation could turn them into inf. + class Model(torch.nn.Module): + def forward(self, a, b): + return torch.hypot(a, b) + + big = torch.finfo(dtype).max + a = torch.tensor([big, big * 0.75, big * 0.7, big * 0.99], dtype=dtype) + b = torch.tensor([1.0, big * 0.6, big * 0.7, big * 0.1], dtype=dtype) + self.assertTrue(torch.isfinite(torch.hypot(a, b)).all()) + + onnx_program = torch.onnx.export(Model(), (a, b), dynamo=True, optimize=False) + _testing.assert_onnx_program(onnx_program) + @parameterized.parameterized.expand( [ ("amax", "amax", False),