Fix

Signed-off-by: Justin Chu <[email protected]>

Author: justinchuby

onnxscript/function_libs/torch_lib/ops/core.py

@@ -933,8 +933,9 @@ def aten_atan2(self: TFloat, other: TFloat) -> TFloat:
     slope = op.Div(self, other)
     atan = op.Atan(slope)
     zero = common_ops.constant(0.0, dtype=self.dtype)
+    pi = common_ops.constant(_MATH_PI, dtype=self.dtype)
 
-    second_third_quadrant = op.Where(op.Greater(self, zero), atan + _MATH_PI, atan - _MATH_PI)
+    second_third_quadrant = op.Where(op.Greater(self, zero), atan + pi, atan - pi)
     result = op.Where(op.Less(other, zero), second_third_quadrant, atan)
 
     # Map NaN to 0 to match PyTorch behavior

tests/function_libs/torch_lib/ops_test_data.py

@@ -578,7 +578,7 @@ def _where_input_wrangler(
     TorchLibOpInfo("asin", core_ops.aten_asin),
     TorchLibOpInfo("asinh", core_ops.aten_asinh),
     TorchLibOpInfo("atan", core_ops.aten_atan),
-    TorchLibOpInfo("atan2", core_ops.aten_atan2, tolerance={torch.float16: (1e-3, 1e-3)}),
+    TorchLibOpInfo("atan2", core_ops.aten_atan2),
     TorchLibOpInfo("atanh", core_ops.aten_atanh),
     TorchLibOpInfo("atleast_1d", core_ops.aten_atleast_1d).skip(
         matcher=lambda sample: isinstance(sample.input, (list, tuple)),