Fix shape inference error

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

Author: justinchuby

onnxscript/optimizer/_constant_folding.py

@@ -919,7 +919,7 @@ def merge_dims(dim1, dim2):
     if other_shape is None:
         return preferred_shape
     if len(preferred_shape) != len(other_shape):
-        raise ValueError("Shapes must have the same rank.")
+        raise ValueError(f"Shapes must have the same rank, got preferred_shape={preferred_shape}, other_shape={other_shape}")
     return ir.Shape(
         [merge_dims(dim1, dim2) for dim1, dim2 in zip(preferred_shape, other_shape)]
     )
@@ -1035,7 +1035,7 @@ def get_type(value: ir.Value) -> onnx.TypeProto | None:
             except Exception as e:
                 logger.debug(
                     "Skipping shape inference for node %r due to exception: %s",
-                    node.name,
+                    node,
                     e,
                 )
 
@@ -1124,7 +1124,12 @@ def process_node(self, node: ir.Node) -> Replacement | None:
         for optimizer in op_optimizers:
             assert optimizer
             context = RewriterContext()
-            output = optimizer(node, context, self._state)
+            try:
+                output = optimizer(node, context, self._state)
+            except Exception as e:
+                raise RuntimeError(
+                    f"Error during constant folding for node {node.name!r} ({node.domain}::{node.op_type})"
+                ) from e
             if output is not None:
                 if isinstance(output, Replacement):
                     return output