diff --git a/onnxscript/function_libs/torch_lib/ops/core.py b/onnxscript/function_libs/torch_lib/ops/core.py index e7f1a611f7..1bb78836e6 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) @@ -4982,10 +5004,42 @@ 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() + 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(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) + 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) + either_inf = op.Or(op.IsInf(self), op.IsInf(other)) + 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 96cd89cf18..a92928b924 100644 --- a/tests/function_libs/torch_lib/e2e_ops_tests.py +++ b/tests/function_libs/torch_lib/e2e_ops_tests.py @@ -2013,6 +2013,90 @@ 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( + [ + ("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( + [ + ("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), diff --git a/tests/function_libs/torch_lib/ops_test_data.py b/tests/function_libs/torch_lib/ops_test_data.py index 15a0cec84c..fa8890b4ef 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), @@ -764,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),