[torchlib] Implement fmax, fmin and hypot

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:

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),

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),