Keep float scalars exact in float64 ops in ScalarToTensorPass (pytorch#24168)

### Summary

Fixes pytorch#24167.

`ScalarToTensorPass` turns a Python scalar operand into
`torch.tensor(value)`, which is float32 for a float. In a float64
program, `RemoveMixedTypeOperators` then casts that constant to float64,
but the value has already been rounded, so `x * 0.1` multiplies by
`0.10000000149011612`. `x + 0.1`, `x - 0.1`, `x / 0.1`, `0.1 - x` and
`torch.deg2rad(x)` were off by about 1e-9 to 1.5e-8 relative.

When the op has a float64 tensor operand, the pass now creates float
scalars as float64 constants. Those graphs also lose the
`_to_dim_order_copy` that cast the constant. Float32, float16 and
integer graphs get the same constants as before.

### Test plan

`test_scalar_to_tensor_pass_keeps_float64_scalars` in
`exir/tests/test_passes.py` runs `(x * 0.1 + 0.3) / 0.7` on a float64
input through `ScalarToTensorPass` and `RemoveMixedTypeOperators` and
checks the result equals eager exactly. It fails before this change and
passes after.

I ran it against executorch 1.5.1 and torch 2.14.1+cpu on Linux x86_64
(WSL2), with this file swapped in (the pass is the same on main). Of the
37 tests in `test_passes.py` matching `scalar`, `mixed`,
`constant_prop`, `to_copy`, `debug_handle`, `symint` or `dim_order`, the
same ones pass before and after apart from the new test. Five fail on
1.5.1 both with and without this change, because main's test file needs
newer `constant_prop_pass` and memory-format code than that release has.
Lowering `x * 0.1`, `x + 0.1`, `x / 0.1`, `0.1 - x` and
`torch.deg2rad(x)` with `to_edge(...).to_executorch()` and running them
through the portable pybindings now matches eager for float64, and gives
the same graph and output as before for float32, float16, and int32
inputs. `ufmt` (black 24.4.2, usort 1.0.8.post1) and `flake8` pass on
the changed files.

This PR was authored with AI assistance (Claude Code).

cc @nil-is-all @JacobSzwejbka @angelayi

Author: usmanmateen

exir/passes/scalar_to_tensor_pass.py

@@ -7,7 +7,7 @@
 # pyre-strict
 
 import torch
-from executorch.exir.pass_base import ExportPass, map_args, PassResult
+from executorch.exir.pass_base import ExportPass, map_args, PassResult, ProxyValue
 
 
 class ScalarToTensorPass(ExportPass):
@@ -17,6 +17,15 @@ def __init__(self) -> None:
 
     # pyre-ignore
     def call_operator(self, op, args, kwargs, meta):
+        # A float scalar next to a float64 tensor must not be rounded to float32
+        # before the op promotes it back to float64.
+        has_float64_tensor = any(
+            isinstance(a, ProxyValue)
+            and a.is_tensor()
+            and a.to_tensor().dtype == torch.float64
+            for a in (*args, *kwargs.values())
+        )
+
         # pyre-ignore
         def try_coerce(value, arg):
             # Note: we want to create tensor constants instead of
@@ -29,7 +38,12 @@ def try_coerce(value, arg):
                     arg.type, torch.TensorType
                 ):
                     self._modified = True
-                    return torch.tensor(value)
+                    dtype = (
+                        torch.float64
+                        if isinstance(value, float) and has_float64_tensor
+                        else None
+                    )
+                    return torch.tensor(value, dtype=dtype)
 
                 return value
 

exir/tests/test_passes.py

@@ -609,6 +609,19 @@ def test_export_scalar_to_tensor_pass(self) -> None:
         new_output = new_gm(test_input)
         self.assertTrue(torch.equal(original_output, new_output))
 
+    def test_scalar_to_tensor_pass_keeps_float64_scalars(self) -> None:
+        class ScalarOps(torch.nn.Module):
+            def forward(self, x: torch.Tensor) -> torch.Tensor:
+                return (x * 0.1 + 0.3) / 0.7
+
+        x = torch.tensor([1.0, 3.0, -7.0, 123456.789], dtype=torch.float64)
+        edge = to_edge(export(ScalarOps(), (x,), strict=True))
+        edge = edge.transform([ScalarToTensorPass(), RemoveMixedTypeOperators()])
+
+        self.assertTrue(
+            torch.equal(edge.exported_program().module()(x), ScalarOps()(x))
+        )
+
     def test_remove_mixed_types_symfloats(self) -> None:
         class Foo(torch.nn.Module):
             def forward(self, x: torch.Tensor) -> torch.Tensor: