Improve Reshape constant fold

Constant fold the input to reshape if from the output shape we know it can be folded

Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>

Author: justinchuby

onnxscript/optimizer/_constant_folding.py

@@ -196,6 +196,10 @@ def set_sym_value(self, value: ir.Value, sym_value: SymbolicValue) -> None:
 
     def get_shape_value(self, value: ir.Value | None) -> ir.Shape | None:
         const_value = _get_numpy_value(value, ir.DataType.INT64, size_limit=10)
+        if const_value is None:
+            # Reshape accepts shape input of INT32 type as well, so we also check for INT32 here
+            # This is common for tflite models
+            const_value = _get_numpy_value(value, ir.DataType.INT32, size_limit=10)
         if const_value is not None:
             if const_value.ndim == 1:
                 return ir.Shape(const_value.tolist())
@@ -471,7 +475,11 @@ def _propagate_shape_value(node: ir.Node, op, state: OptimizerState) -> ReturnVa
 def reshape(node: ir.Node, op, state: OptimizerState) -> ReturnValue:
     """Replace a Reshape node by Identity when applicable.
 
-    Also propagate symbolic shape values.
+    Also propagate symbolic shape values.  When the output shape is known
+    (from shape inference) and the shape input is not already a constant,
+    replace the shape input with a concrete constant:
+    - Fully static output shape → constant with exact dims.
+    - Exactly one symbolic dim → replace it with -1 (Reshape infers it).
     """
     input = _get_input(node, 0)
     shape = _get_input(node, 1)
@@ -482,6 +490,10 @@ def reshape(node: ir.Node, op, state: OptimizerState) -> ReturnValue:
     shape_value = state.get_shape_value(shape)
 
     if shape_value is None or input_shape is None:
+        # Try to materialize the shape input from the output shape
+        result = _try_materialize_reshape_shape(node, op, input, shape)
+        if result is not None:
+            return result
         return _propagate_shape_value(node, op, state)
 
     # No need to check for special values like -1, 0, etc. here
@@ -490,6 +502,44 @@ def reshape(node: ir.Node, op, state: OptimizerState) -> ReturnValue:
     return _propagate_shape_value(node, op, state)
 
 
+def _try_materialize_reshape_shape(
+    node: ir.Node, op, input: ir.Value, shape: ir.Value
+) -> ReturnValue:
+    """Create a constant shape input for Reshape from the output shape.
+
+    When the Reshape output already has a known shape (e.g., from symbolic
+    shape inference) but the shape input is not a constant, we can replace
+    the shape input with a materialized constant tensor.
+
+    - If the output shape is fully static, use the exact dims.
+    - If exactly one dim is symbolic, replace it with -1.
+    """
+    output = _get_output(node, 0)
+    if output is None or output.shape is None:
+        return None
+
+    out_shape = output.shape
+    dims = list(out_shape)
+    sym_count = sum(1 for d in dims if not isinstance(d, int))
+
+    if sym_count == 0:
+        # Fully static — create constant with exact dims
+        new_shape = op.Constant(
+            value=ir.tensor(dims, dtype=ir.DataType.INT64),
+        )
+        return op.Reshape(input, new_shape)
+
+    if sym_count == 1:
+        # Replace the single symbolic dim with -1
+        concrete_dims = [-1 if not isinstance(d, int) else d for d in dims]
+        new_shape = op.Constant(
+            value=ir.tensor(concrete_dims, dtype=ir.DataType.INT64),
+        )
+        return op.Reshape(input, new_shape)
+
+    return None
+
+
 @register("Squeeze")
 def squeeze(node: ir.Node, op, state: OptimizerState) -> ReturnValue:
     """Propagate symbolic shape values."""