Repository: microsoft/onnxscript

Issue:
llama with dort is failing
dort is failing with torch-nightly, onnx (development version), transformers==4.37.2

To replicate, run test ``test_llama_mixed_precision_large`` in https://github.com/microsoft/onnxscript/pull/1317.

```
  File "/usr/local/lib/python3.10/dist-packages/torch/fx/graph_module.py", line 735, in call_wrapped
    return self._wrapped_call(self, *args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/onnx/_internal/onnxruntime.py", line 931, in _ort_acclerated_call
    onnx_model = optimizer.optimize(onnx_model)
  File "/src/github/onnxscript/onnxscript/optimizer/__init__.py", line 61, in optimize
    model = onnx.shape_inference.infer_shapes(
  File "/usr/local/lib/python3.10/dist-packages/onnx/shape_inference.py", line 46, in infer_shapes
    inferred_model_str = C.infer_shapes(
onnx.onnx_cpp2py_export.shape_inference.InferenceError: [ShapeInferenceError] Inference error(s): (op_type:ConstantOfShape, node name: node_ConstantOfShape_3): [TypeInferenceError] Attribute expected to have a one-dim tensor
(op_type:ConstantOfShape, node name: node_ConstantOfShape_2): [TypeInferenceError] Attribute expected to have a one-dim tensor
(op_type:Transpose, node name: Transpose_408): [TypeInferenceError] Input 0 expected to have type but instead is null
(op_type:ConstantOfShape, node name: node_ConstantOfShape_1): [TypeInferenceError] Attribute expected to have a one-dim tensor
(op_type:Transpose, node name: Transpose_431): [TypeInferenceError] Input 0 expected to have type but instead is null
(op_type:ConstantOfShape, node name: node_ConstantOfShape_0): [TypeInferenceError] Attribute expected to have a one-dim tensor
(op_type:Transpose, node name: Transpose_453): [TypeInferenceError] Input 0 expected to have type but instead is null
(op_type:Transpose, node name: Transpose_474): [TypeInferenceError] Input 0 expected to have type but instead is null
```



Discussion:

Resolving patch:
diff --git a/onnxscript/rewriter/cast_constant_of_shape.py b/onnxscript/rewriter/cast_constant_of_shape.py
--- a/onnxscript/rewriter/cast_constant_of_shape.py
+++ b/onnxscript/rewriter/cast_constant_of_shape.py
@@ -19,7 +19,7 @@ def cast_constant_of_shape(shape, scalar, dtype):
 def fused_cast_constant_of_shape(op, shape: ir.Value, scalar: ir.Attr, dtype: ir.Attr, **_):
     # Cast scalar (a TensorProto attribute) to the specified dtype
     scalar_value = scalar.value.numpy().item()
-    cast_value = onnx.helper.make_tensor("value", dtype.value, (), [scalar_value])
+    cast_value = onnx.helper.make_tensor("value", dtype.value, (1,), [scalar_value])
     return op.ConstantOfShape(shape, value=cast_value)
 
 
@@ -29,7 +29,7 @@ def cast_constant_of_shape_without_value(shape, dtype):
 
 
 def fused_cast_constant_of_shape_without_value(op, shape, dtype, **_):
-    zero = onnx.helper.make_tensor("value", dtype.value, (), [0])
+    zero = onnx.helper.make_tensor("value", dtype.value, (1,), [0])
     return op.ConstantOfShape(shape, value=zero)
 
 
diff --git a/onnxscript/rewriter/pattern.py b/onnxscript/rewriter/pattern.py
--- a/onnxscript/rewriter/pattern.py
+++ b/onnxscript/rewriter/pattern.py
@@ -221,36 +221,20 @@ class OpsetPattern:
     of OpPattern, and  `op.Matmul(x, y)` is an instance of NodePattern.
 
     An opset pattern is also matched against the actual opset used in the
-    input model. Typically, we match against an ONNX opset (ignoring the
-    version), but we can match against a specific version of the opset too.
-    However, it is preferable that version-dependences are handled at the
-    level of a rewrite rule, rather than at the level of a pattern.
-
+    input model.
     """
 
-    def __init__(
-        self,
-        domain_pattern: PythonPattern | PrefixPattern,
-        version_pattern: PythonPattern | AnyPattern,
-    ) -> None:
+    def __init__(self, domain_pattern: PythonPattern | PrefixPattern | str) -> None:
+        if isinstance(domain_pattern, str):
+            domain_pattern = PythonPattern(domain_pattern)
         self.domain_pattern = domain_pattern
-        self.version_pattern = version_pattern
-
-    @classmethod
-    def singleton(cls, domain: str, version: int) -> OpsetPattern:
-        return cls(PythonPattern(domain), PythonPattern(version))
-
-    @classmethod
-    def domain(cls, domain: str) -> OpsetPattern:
-        return cls(PythonPattern(domain), AnyPattern())
 
     @classmethod
     def domain_prefix(cls, domain: str) -> OpsetPattern:
-        return cls(PrefixPattern(domain), AnyPattern())
+        return cls(PrefixPattern(domain))
 
-    def matches(self, opset):
-        domain, version = opset
-        return self.domain_pattern.matches(domain) and self.version_pattern.matches(version)
+    def matches(self, domain):
+        return self.domain_pattern.matches(domain)
 
     def __getattr__(self, name: str) -> Any:
         return OpPattern(self, PythonPattern(name))
@@ -260,11 +244,9 @@ def submodule(self, name: str) -> Any:
         return OpPattern(self, PrefixPattern(name))
 
 
-opset17 = OpsetPattern.singleton("", 17)
-
-onnxop = OpsetPattern.domain("")
+onnxop = OpsetPattern("")
 
-msft_op = OpsetPattern.singleton("com.microsoft", 1)
+msft_op = OpsetPattern("com.microsoft")
 
 torch_module_op = OpsetPattern.domain_prefix("pkg.torch")
 
@@ -383,8 +365,7 @@ def fail(self):
         self.matched_values = None
         self.bindings = {}
 
-    def extend(self, other: MatchResult | bool, model):
-        del model  # Unused
+    def extend(self, other: MatchResult | bool):
         if not self.success:
             return
         if not other:
@@ -423,6 +404,12 @@ def matches(self, value: ir.Value, model: ir.Model):
         return MatchResult([], {self.name: value})
 
     def commute(self) -> Sequence[ValuePattern]:
+        """Return a list of commuted patterns.
+
+        This is used to handle commutative operations like addition and multiplication.
+        A single pattern is converted into a list of equivalent patterns by swapping
+        the parameters of commutative operations.
+        """
         return [self]
 
     def __add__(self, other):
@@ -474,24 +461,9 @@ def __init__(
         self.attributes = attributes
         self.bound_value = None
 
-    def matches(self, value: ir.Value, model: ir.Model):
-        if self.bound_value is not None:
-            # DAG-matching, not Tree-matching.
-            if self.bound_value.is_same_as(value):
-                return MatchResult([])
-            else:
-                return MatchResult.FAIL()
-        node = value.producer()
-        if node is None:
-            # Eg., value could be an input parameter, which will not match a value
-            # computed by the op in this pattern.
-            return MatchResult.FAIL()
-        return self.matches_node(node, model)
-
     def matches_node(self, node: ir.Node, model: ir.Model) -> MatchResult:
         """Examine if the IR node matches the self pattern."""
-        node_version = model.graph.opset_imports.get(node.domain, 1)
-        if not self.domain.matches((node.domain, node_version)):
+        if not self.domain.matches(node.domain):
             return MatchResult.FAIL()
         if not self.op.matches(node.op_type):
             return MatchResult.FAIL()
@@ -506,7 +478,7 @@ def matches_node(self, node: ir.Node, model: ir.Model) -> MatchResult:
             if arg_value is None or previous_node_output_pattern is None:
                 return MatchResult.FAIL()
             sub_match = previous_node_output_pattern.matches(arg_value, model)  # type: ignore[attr-defined]
-            match.extend(sub_match, model)
+            match.extend(sub_match)
             if not match:  # If sub-match failed,
                 return match
         # Sub-graphs not handled yet.
@@ -517,7 +489,7 @@ def matches_node(self, node: ir.Node, model: ir.Model) -> MatchResult:
             sub_match = attr_pattern.matches(attr_value, model)  # type: ignore[arg-type]
             if not sub_match:
                 return MatchResult.FAIL()
-            match.extend(sub_match, model)
+            match.extend(sub_match)
         for name in node.attributes:
             # TODO: Support matching default values for attributes.
             if name not in self.attributes:
@@ -540,9 +512,7 @@ def enumerate_inputs(inputs, index):
                         yield [pattern, *rest]
 
         inputs = list(enumerate_inputs(list_of_lists, 0))
-        if self.domain.matches(("", None)) and (
-            self.op.matches("Add") or self.op.matches("Mul")
-        ):
+        if self.domain.matches("") and (self.op.matches("Add") or self.op.matches("Mul")):
             # TODO: handle cases where number of inputs is not 2.
             swapped = [[x[1], x[0]] for x in inputs]
             inputs.extend(swapped)
@@ -572,6 +542,12 @@ def matches(self, value: ir.Value, model: ir.Model):
             return MatchResult.FAIL()
         return self.node_pattern.matches_node(node, model)
 
+    def commute(self) -> Sequence[ValuePattern]:
+        return [
+            NodeOutputPattern(pattern, self.output_index, self.name)
+            for pattern in self.node_pattern.commute()
+        ]
+
 
 Var = ValuePattern
 
@@ -649,8 +625,11 @@ def pattern(x, shape1, shape2):
         num_outputs = len(node_output_pattern)
         for i, p in enumerate(node_output_pattern):
             assert isinstance(p, NodeOutputPattern)
-            assert p.node_pattern is node_pattern
-            assert p.output_index == i
+            if (p.node_pattern is not node_pattern) or (p.output_index != i):
+                raise NotImplementedError(
+                    "Multi-output pattern not handled by this API. "
+                    "Use other APIs to handle multi-output patterns."
+                )
     else:
         raise TypeError(f"Invalid type {type(node_output_pattern)} for pattern")
     return node_pattern, num_outputs


Test patch:
diff --git a/onnxscript/rewriter/cast_constant_of_shape_test.py b/onnxscript/rewriter/cast_constant_of_shape_test.py
--- a/onnxscript/rewriter/cast_constant_of_shape_test.py
+++ b/onnxscript/rewriter/cast_constant_of_shape_test.py
@@ -1,6 +1,8 @@
 import unittest
 
+import onnx.checker
 import onnx.parser
+import onnx.printer
 
 from onnxscript import ir
 from onnxscript.rewriter import cast_constant_of_shape
@@ -8,21 +10,27 @@
 
 class CastConstantOfShapeTest(unittest.TestCase):
     def test_cast_after_constant_of_shape_is_fused(self):
-        model_proto = onnx.parser.parse_model(
+        input_model_proto = onnx.parser.parse_model(
             """
             <ir_version: 7, opset_import: [ "" : 17]>
             agraph (int64[2] input_x) => (float16[1, 4] output)
             {
                 constant = ConstantOfShape <value: tensor = float[1] {1.}>(input_x)
-                output = Cast <to = 10> (constant)
+                temp = Cast <to = 10> (constant)
+                output = Identity (temp)
             }
             """
         )
-        model = ir.serde.deserialize_model(model_proto)
+        onnx.checker.check_model(input_model_proto, True)
+        model = ir.serde.deserialize_model(input_model_proto)
         count = cast_constant_of_shape.rules.apply_to_model(model)
         self.assertEqual(count, 1)
-        self.assertEqual(len(model.graph), 1)
+        self.assertEqual(len(model.graph), 2)
         self.assertEqual(model.graph[0].attributes["value"].value.dtype, 10)
+        output_model_proto = ir.serde.serialize_model(model)
+        # TODO: Eliminating `temp` in above example causes a failure.
+        # Rewriter changes graph output name, doesn't introduce type for it
+        onnx.checker.check_model(output_model_proto, True)
 
     def test_cast_after_constant_of_shape_without_value_is_fused(self):
         model_proto = onnx.parser.parse_model(
@@ -31,15 +39,18 @@ def test_cast_after_constant_of_shape_without_value_is_fused(self):
             agraph (int64[2] input_x) => (float16[1, 4] output)
             {
                 constant = ConstantOfShape (input_x)
-                output = Cast <to = 10> (constant)
+                temp = Cast <to = 10> (constant)
+                output = Identity (temp)
             }
             """
         )
         model = ir.serde.deserialize_model(model_proto)
         count = cast_constant_of_shape.rules.apply_to_model(model)
         self.assertEqual(count, 1)
-        self.assertEqual(len(model.graph), 1)
+        self.assertEqual(len(model.graph), 2)
         self.assertEqual(model.graph[0].attributes["value"].value.dtype, 10)
+        output_model_proto = ir.serde.serialize_model(model)
+        onnx.checker.check_model(output_model_proto, True)
 
 
 if __name__ == "__main__":