Fix cast inference bug (#1884)

Fix bug introduced while migrating cast-logic to new IR in optimizer.

---------

Signed-off-by: gramalingam <[email protected]>

Author: gramalingam

onnxscript/ir/_core.py

@@ -879,6 +879,10 @@ def __init__(
             )
         self._frozen: bool = frozen
 
+    def copy(self):
+        """Return a copy of the shape."""
+        return Shape(self._dims, self._denotations, self._frozen)
+
     @property
     def dims(self) -> tuple[int | SymbolicDim, ...]:
         """All dimensions in the shape.

onnxscript/optimizer/_constant_folding.py

@@ -236,7 +236,13 @@ def cast(node: ir.Node, op, state: OptimizerState) -> ReturnValue:
     input = _get_input(node, 0)
     output = _get_output(node, 0)
     if input is not None and output is not None:
-        _update_type(output, input.type)
+        input_shape = input.shape
+        if input_shape is not None:
+            output.shape = input_shape.copy()
+    if output is not None:
+        output_dtype = _get_int_attribute(node, "to", None)
+        if output_dtype is not None:
+            output.type = ir.TensorType(ir.DataType(output_dtype))
     return None