Fix infinite loop in constant folding of Resize: pre-evaluate output size check

Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>

Author: Copilot

onnxscript/optimizer/_constant_folding.py

@@ -1280,6 +1280,31 @@ def process_node(self, node: ir.Node, is_function: bool) -> Replacement | None:
                 node.name,
             )
 
+        # Check estimated output size before calling the (potentially expensive) reference
+        # evaluator. This avoids hanging on ops like Resize with large output tensors where
+        # the reference implementation may be very slow.
+        for output in node.outputs:
+            if output.shape is not None and output.shape.is_static():
+                try:
+                    shape_tuple = output.shape.numpy()
+                    estimated_size = math.prod(shape_tuple)
+                    if estimated_size > self.output_size_limit:
+                        logger.info(
+                            "Skipping constant folding for node %r due to large estimated "
+                            "output size: %s > output_size_limit=%s",
+                            node.name,
+                            estimated_size,
+                            self.output_size_limit,
+                        )
+                        return None
+                except Exception as e:  # pylint: disable=broad-exception-caught
+                    logger.debug(
+                        "Could not estimate output size for node %r output '%s': %s",
+                        node.name,
+                        output.name,
+                        e,
+                    )
+
         input_values = [_get_numpy_value(x) for x in node.inputs]
 
         def convert(av):

onnxscript/optimizer/_constant_folding_test.py

@@ -741,6 +741,38 @@ def test_constant_folding_creates_constant_nodes_in_function(self):
             constant_nodes = [n for n in func.graph if n.op_type == "Constant"]
             self.assertEqual(len(constant_nodes), 1)
 
+    def test_output_size_limit_prevents_evaluation_before_running(self):
+        """Test that output_size_limit is checked before calling the (expensive) reference
+        evaluator. This is a regression test for https://github.com/microsoft/onnxscript/issues/2709
+        where Resize with cubic mode on a large output tensor caused an infinite loop/hang
+        because the reference evaluator was called before checking output size.
+        """
+        # Create a model with Resize that has a small constant input but large output
+        # sizes: [1, 2, 64, 64] -> output has 8192 elements, exceeding a small output_size_limit
+        model_text = """
+            <ir_version: 8, opset_import: [ "" : 18]>
+            agraph () => (float[1, 2, 64, 64] output)
+            <float[1, 2, 4, 4] x_const = {1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0,
+                                            9.0, 10.0, 11.0, 12.0, 13.0, 14.0, 15.0, 16.0,
+                                            1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0,
+                                            9.0, 10.0, 11.0, 12.0, 13.0, 14.0, 15.0, 16.0}>
+            {
+                sizes = Constant <value_ints=[1, 2, 64, 64]> ()
+                output = Resize <mode="nearest"> (x_const, , , sizes)
+            }
+        """
+        model = ir.from_onnx_text(model_text)
+        # With a small output_size_limit, the Resize node should NOT be folded
+        # (and more importantly, should NOT hang)
+        optimized = self._fold(
+            model,
+            onnx_shape_inference=True,
+            output_size_limit=100,  # very small limit to prevent evaluation
+        )
+        # The Resize node should remain (not folded)
+        ops = [node.op_type for node in optimized.graph]
+        self.assertIn("Resize", ops)
+
     def test_initializer_as_graph_output_is_not_removed(self):
         """Test that an initializer that is a graph output is not removed during constant folding."""
         model = """