Keep float64 precision of Python floats in scalar_tensor

onnxscript/function_libs/torch_lib/ops/core.py

@@ -9085,6 +9085,10 @@ def aten_scalar_tensor(
     if dtype == -1:
         dtype = FLOAT.dtype
 
+    if isinstance(s, float) and dtype in (DOUBLE.dtype, COMPLEX128.dtype):
+        # Cast would make a float32 constant from the Python float and round 0.1 to
+        # 0.10000000149011612.
+        s = op.Constant(value=ir.tensor(s, dtype=ir.DataType.DOUBLE))
     return common_ops.cast_to(s, dtype=dtype)
 
 

tests/function_libs/torch_lib/e2e_ops_tests.py

@@ -1967,6 +1967,19 @@ def forward(self, x):
                 self.assertEqual(node.attributes.get_int("noop_with_empty_axes", 0), 0)
         _testing.assert_onnx_program(onnx_program)
 
+    def test_python_float_keeps_float64_precision(self):
+        # torch.export turns the Python floats into scalar_tensor(0.1, dtype=float64).
+        class Model(torch.nn.Module):
+            def forward(self, x):
+                return x + 0.1, x * 0.1, x == 0.1, torch.clamp(x, max=0.1)
+
+        x = torch.tensor([1.0, 0.3, 0.1, -2.7], dtype=torch.float64)
+        onnx_program = torch.onnx.export(Model(), (x,), dynamo=True, verbose=False)
+        # assert_onnx_program allows 1e-7 for float64, more than 0.1 rounded to float32
+        # is off by, so compare exactly.
+        for actual, expected in zip(onnx_program(x), Model()(x)):
+            torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
 
 if __name__ == "__main__":
     unittest.main()