Repository: microsoft/onnxscript
Pull request #2175: Introduce pattern.any_value
Author: login: gramalingam
Labels: name: module: rewriter
Created: 2025-04-09T22:14:27Z | Merged: 2025-04-10T20:05:57Z

Introduce pattern.any_value as a convenience when writing patterns. It is more precise than using `_allow_other_inputs=True` (which will allow any number of inputs).



Diff:
diff --git a/onnxscript/rewriter/ort_fusions/gqa.py b/onnxscript/rewriter/ort_fusions/gqa.py
index 7f761a3744..266987dd4d 100644
--- a/onnxscript/rewriter/ort_fusions/gqa.py
+++ b/onnxscript/rewriter/ort_fusions/gqa.py
@@ -87,19 +87,17 @@ def pattern(
         shape_B111,
     ):
         # Reshape query from (B, S, D) to (B, S, H, D/H)
-        query_BSHDh = op.Reshape(query_BSD, _allow_other_inputs=True, _outputs=["query_BSHDh"])
+        query_BSHDh = op.Reshape(query_BSD, pattern.ANY_VALUE, _outputs=["query_BSHDh"])
         # Transpose from (B, S, H, D/H) to (B, H, S, D/H)
         query_BHSDh = op.Transpose(query_BSHDh, perm=[0, 2, 1, 3])
 
         # Reshape key from (B, S, Dkv) to (B, S, Hkv, D/H)
-        key_BSHkvDh = op.Reshape(key_BSDkv, _allow_other_inputs=True, _outputs=["key_BSHkvDh"])
+        key_BSHkvDh = op.Reshape(key_BSDkv, pattern.ANY_VALUE, _outputs=["key_BSHkvDh"])
         # Transpose from (B, S, Hkv, D/H) to (B, Hkv, S, D/H)
         key_BHkvSDh = op.Transpose(key_BSHkvDh, perm=[0, 2, 1, 3])
 
         # Reshape value from (B, S, Dkv) to (B, S, Hkv, D/H)
-        value_BSHkvDh = op.Reshape(
-            value_BSDkv, _allow_other_inputs=True, _outputs=["value_BSHkvDh"]
-        )
+        value_BSHkvDh = op.Reshape(value_BSDkv, pattern.ANY_VALUE, _outputs=["value_BSHkvDh"])
         # Transpose from (B, S, Hkv, D/H) to (B, Hkv, S, D/H)
         value_BHkvSDh = op.Transpose(value_BSHkvDh, perm=[0, 2, 1, 3])
 
@@ -129,18 +127,18 @@ def pattern(
 
         key_seq_BHkvTDh = op.Concat(past_key, key_BHkvSDh_rope, axis=-2)
         key_seq_BHkv1TDh = op.Unsqueeze(key_seq_BHkvTDh, 2)
-        key_seq_BHkvGTDh = op.Expand(key_seq_BHkv1TDh, _allow_other_inputs=True)
+        key_seq_BHkvGTDh = op.Expand(key_seq_BHkv1TDh, pattern.ANY_VALUE)
         key_seq_BHTDh = op.Reshape(
-            key_seq_BHkvGTDh, _allow_other_inputs=True, _outputs=["key_seq_BHTDh"]
+            key_seq_BHkvGTDh, pattern.ANY_VALUE, _outputs=["key_seq_BHTDh"]
         )
 
         # Concatenate past_value cache and current value, expand across heads
         # that share key/value.
         value_seq_BHkvTDh = op.Concat(past_value, value_BHkvSDh, axis=-2)
         value_seq_BHkv1TDh = op.Unsqueeze(value_seq_BHkvTDh, 2)
-        value_seq_BHkvGTDh = op.Expand(value_seq_BHkv1TDh, _allow_other_inputs=True)
+        value_seq_BHkvGTDh = op.Expand(value_seq_BHkv1TDh, pattern.ANY_VALUE)
         value_seq_BHTDh = op.Reshape(
-            value_seq_BHkvGTDh, _allow_other_inputs=True, _outputs=["value_seq_BHTDh"]
+            value_seq_BHkvGTDh, pattern.ANY_VALUE, _outputs=["value_seq_BHTDh"]
         )
 
         mask = causal_mask_pattern(op, input_ids, some_kv_cache, shape_B111)
@@ -158,7 +156,7 @@ def pattern(
         attention_BSHDh = op.Transpose(attention_BHSDh, perm=[0, 2, 1, 3])
         # Reshape back to (B, S, D)
         attention_BSD = op.Reshape(
-            attention_BSHDh, _allow_other_inputs=True, _outputs=["attention_BSD"]
+            attention_BSHDh, pattern.ANY_VALUE, _outputs=["attention_BSD"]
         )
         return attention_BSD, key_seq_BHkvTDh, value_seq_BHkvTDh
 
diff --git a/onnxscript/rewriter/pattern.py b/onnxscript/rewriter/pattern.py
index 907ebd0b88..cfca31125f 100644
--- a/onnxscript/rewriter/pattern.py
+++ b/onnxscript/rewriter/pattern.py
@@ -634,6 +634,20 @@ def _is_pattern_variable(x: Any) -> bool:
     return type(x) is ValuePattern
 
 
+class AnyValue(ValuePattern):
+    """Represents a pattern that matches against any value."""
+
+    def __init__(self) -> None:
+        super().__init__(None)
+
+    def clone(self, node_map: dict[NodePattern, NodePattern]) -> AnyValue:
+        # A single instance of AnyValue suffices.
+        return self
+
+
+ANY_VALUE = AnyValue()
+
+
 class Constant(ValuePattern):
     """Represents a pattern that matches against a scalar constant value."""
 
@@ -1108,6 +1122,9 @@ def _bind_value(self, pattern_value: ValuePattern, value: ir.Value | None) -> bo
 
     def _match_value(self, pattern_value: ValuePattern, value: ir.Value | None) -> bool:
         """Match an IR value against a ValuePattern instance."""
+        if isinstance(pattern_value, AnyValue):
+            return True
+
         if not self._bind_value(pattern_value, value):
             return False
 
diff --git a/onnxscript/rewriter/pattern_test.py b/onnxscript/rewriter/pattern_test.py
index 24ae237c20..ce11e23c19 100644
--- a/onnxscript/rewriter/pattern_test.py
+++ b/onnxscript/rewriter/pattern_test.py
@@ -667,6 +667,27 @@ def test_model(x: FLOAT[1024], y: FLOAT[1024], z: FLOAT[1024]) -> FLOAT[1024]:
         onnxscript.optimizer.inline(model)
         self.assertEqual([x.op_type for x in model.graph], ["Add", "Mul", "Add", "Mul"])
 
+    def test_any_value(self):
+        def source_pattern(op, x):
+            return op.Add(x, op.Mul(0, pattern.ANY_VALUE))
+
+        def replacement(op, x):
+            return op.Identity(x)
+
+        rule = pattern.RewriteRule(source_pattern, replacement)
+
+        @script()
+        def test_model(x: FLOAT[1024], y: FLOAT[1024]) -> FLOAT[1024]:
+            zero = op.Constant(value_float=0.0)
+            return op.Add(x, op.Mul(zero, y))
+
+        model_proto = test_model.to_model_proto()
+        model = ir.serde.deserialize_model(model_proto)
+        self.assertEqual([x.op_type for x in model.graph], ["Constant", "Mul", "Add"])
+        rule.apply_to_model(model)
+        self.assertEqual(len(model.graph), 2)
+        self.assertEqual([x.op_type for x in model.graph], ["Constant", "Identity"])
+
 
 class PatternBuilderTest(unittest.TestCase):
     def test_pattern_builder_context(self):


Reviews:
user: login: gramalingam
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-09T22:15:51Z
user: login: justinchuby
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-09T22:45:46Z
user: login: shubhambhokare1
body: Could this be extended for _allow_other_attributes as well?
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-10T06:14:09Z
user: login: shubhambhokare1
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-10T06:14:41Z
user: login: gramalingam
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-10T16:11:49Z
user: login: justinchuby
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-04-10T16:23:53Z
user: login: justinchuby
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-04-10T19:55:27Z

Review comments:
diff_hunk: @@ -634,6 +634,20 @@ def _is_pattern_variable(x: Any) -> bool:
     return type(x) is ValuePattern
 
 
+class AnyValue(ValuePattern):
+    """Represents a pattern that matches against any value."""
+
+    def __init__(self) -> None:
+        super().__init__(None)
+
+    def clone(self, node_map: dict[NodePattern, NodePattern]) -> AnyValue:
+        # A single instance of AnyValue suffices.
+        return self
+
+
+any_value = AnyValue()
path: onnxscript/rewriter/pattern.py
user: login: gramalingam
body: Any suggestions for alternative names? Another possibility is `ignore`. Or, would altering the casing to `ANY_VALUE` be better?
created_at: 2025-04-09T22:15:50Z
author_association: COLLABORATOR
diff_hunk: @@ -634,6 +634,20 @@ def _is_pattern_variable(x: Any) -> bool:
     return type(x) is ValuePattern
 
 
+class AnyValue(ValuePattern):
+    """Represents a pattern that matches against any value."""
+
+    def __init__(self) -> None:
+        super().__init__(None)
+
+    def clone(self, node_map: dict[NodePattern, NodePattern]) -> AnyValue:
+        # A single instance of AnyValue suffices.
+        return self
+
+
+any_value = AnyValue()
path: onnxscript/rewriter/pattern.py
user: login: justinchuby
body: I would vote for `ANY_VALUE`
created_at: 2025-04-09T22:45:46Z
author_association: COLLABORATOR
diff_hunk: @@ -634,6 +634,20 @@ def _is_pattern_variable(x: Any) -> bool:
     return type(x) is ValuePattern
 
 
+class AnyValue(ValuePattern):
+    """Represents a pattern that matches against any value."""
+
+    def __init__(self) -> None:
+        super().__init__(None)
+
+    def clone(self, node_map: dict[NodePattern, NodePattern]) -> AnyValue:
+        # A single instance of AnyValue suffices.
+        return self
+
+
+any_value = AnyValue()
path: onnxscript/rewriter/pattern.py
user: login: shubhambhokare1
body: Leaning towards `ANY_VALUE` as well
created_at: 2025-04-10T06:14:41Z
author_association: CONTRIBUTOR
diff_hunk: @@ -634,6 +634,20 @@ def _is_pattern_variable(x: Any) -> bool:
     return type(x) is ValuePattern
 
 
+class AnyValue(ValuePattern):
+    """Represents a pattern that matches against any value."""
+
+    def __init__(self) -> None:
+        super().__init__(None)
+
+    def clone(self, node_map: dict[NodePattern, NodePattern]) -> AnyValue:
+        # A single instance of AnyValue suffices.
+        return self
+
+
+any_value = AnyValue()
path: onnxscript/rewriter/pattern.py
user: login: gramalingam
body: done
created_at: 2025-04-10T16:11:49Z
author_association: COLLABORATOR

PR comments:
user: login: codecov[bot]
created_at: 2025-04-09T22:18:33Z
body: ### :x: 4 Tests Failed:
| Tests completed | Failed | Passed | Skipped |
|---|---|---|---|
| 14115 | 4 | 14111 | 1697 |
<details><summary>View the top 3 failed test(s) by shortest run time</summary>

> 
> ```python
> onnxscript.backend.onnx_export_test.TestOnnxBackEnd::test_export2python_produces_correct_onnx_script_model_0564_test_log
> ```
> 
> <details><summary>Stack Traces | 0.003s run time</summary>
> 
> > 
> > ```python
> > onnxscript\backend\onnx_export_test.py:137: in extract_functions
> >     mod = importlib.import_module(import_name)
> > C:\hostedtoolcache\windows\Python\3.10.11\x64\lib\importlib\__init__.py:126: in import_module
> >     return _bootstrap._gcd_import(name[level:], package, level)
> > E   ModuleNotFoundError: No module named 'tests.onnx_backend_test_code.test_log'
> > 
> > The above exception was the direct cause of the following exception:
> > .nox\test\lib\site-packages\parameterized\parameterized.py:620: in standalone_func
> >     return func(*(a + p.args), **p.kwargs, **kw)
> > onnxscript\backend\onnx_export_test.py:271: in test_export2python_produces_correct_onnx_script_model
> >     functions = extract_functions(backend_test.name, code, self.test_folder)
> > onnxscript\backend\onnx_export_test.py:139: in extract_functions
> >     raise AssertionError(
> > E   AssertionError: Unable to import 'tests.onnx_backend_test_code.test_log' (e=No module named 'tests.onnx_backend_test_code.test_log') (file: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_log.py', absolute path: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_log.py', current folder: D:\a\onnxscript\onnxscript
> > E   ---- CONTENT --
> > E   import numpy
> > E   from onnx import TensorProto
> > E   from onnx.helper import make_tensor
> > E   from onnxscript import script, external_tensor
> > E   from onnxscript.values import Opset
> > E   from onnxscript.onnx_types import FLOAT
> > E   from onnxscript.onnx_opset import opset13
> > E   
> > E   @script()
> > E   def bck_test_log(x: FLOAT[3,4,5]) -> (FLOAT[3,4,5]):
> > E       y = opset13.Log(x)
> > E       return y
> > ```
> 
> </details>


> 
> ```python
> onnxscript.backend.onnx_export_test.TestOnnxBackEnd::test_export2python_produces_correct_onnx_script_model_1134_test_softmax_axis_1
> ```
> 
> <details><summary>Stack Traces | 0.003s run time</summary>
> 
> > 
> > ```python
> > onnxscript\backend\onnx_export_test.py:137: in extract_functions
> >     mod = importlib.import_module(import_name)
> > C:\hostedtoolcache\windows\Python\3.10.11\x64\lib\importlib\__init__.py:126: in import_module
> >     return _bootstrap._gcd_import(name[level:], package, level)
> > E   ModuleNotFoundError: No module named 'tests.onnx_backend_test_code.test_softmax_axis_1'
> > 
> > The above exception was the direct cause of the following exception:
> > .nox\test\lib\site-packages\parameterized\parameterized.py:620: in standalone_func
> >     return func(*(a + p.args), **p.kwargs, **kw)
> > onnxscript\backend\onnx_export_test.py:271: in test_export2python_produces_correct_onnx_script_model
> >     functions = extract_functions(backend_test.name, code, self.test_folder)
> > onnxscript\backend\onnx_export_test.py:139: in extract_functions
> >     raise AssertionError(
> > E   AssertionError: Unable to import 'tests.onnx_backend_test_code.test_softmax_axis_1' (e=No module named 'tests.onnx_backend_test_code.test_softmax_axis_1') (file: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_softmax_axis_1.py', absolute path: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_softmax_axis_1.py', current folder: D:\a\onnxscript\onnxscript
> > E   ---- CONTENT --
> > E   import numpy
> > E   from onnx import TensorProto
> > E   from onnx.helper import make_tensor
> > E   from onnxscript import script, external_tensor
> > E   from onnxscript.values import Opset
> > E   from onnxscript.onnx_types import FLOAT
> > E   from onnxscript.onnx_opset import opset13
> > E   
> > E   @script()
> > E   def bck_test_softmax_axis_1(x: FLOAT[3,4,5]) -> (FLOAT[3,4,5]):
> > E       y = opset13.Softmax(x, axis=1)
> > E       return y
> > ```
> 
> </details>


> 
> ```python
> onnxscript.backend.onnx_export_test.TestOnnxBackEnd::test_export2python_produces_correct_onnx_script_model_1143_test_softmax_example
> ```
> 
> <details><summary>Stack Traces | 0.003s run time</summary>
> 
> > 
> > ```python
> > onnxscript\backend\onnx_export_test.py:137: in extract_functions
> >     mod = importlib.import_module(import_name)
> > C:\hostedtoolcache\windows\Python\3.10.11\x64\lib\importlib\__init__.py:126: in import_module
> >     return _bootstrap._gcd_import(name[level:], package, level)
> > E   ModuleNotFoundError: No module named 'tests.onnx_backend_test_code.test_softmax_example'
> > 
> > The above exception was the direct cause of the following exception:
> > .nox\test\lib\site-packages\parameterized\parameterized.py:620: in standalone_func
> >     return func(*(a + p.args), **p.kwargs, **kw)
> > onnxscript\backend\onnx_export_test.py:271: in test_export2python_produces_correct_onnx_script_model
> >     functions = extract_functions(backend_test.name, code, self.test_folder)
> > onnxscript\backend\onnx_export_test.py:139: in extract_functions
> >     raise AssertionError(
> > E   AssertionError: Unable to import 'tests.onnx_backend_test_code.test_softmax_example' (e=No module named 'tests.onnx_backend_test_code.test_softmax_example') (file: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_softmax_example.py', absolute path: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_softmax_example.py', current folder: D:\a\onnxscript\onnxscript
> > E   ---- CONTENT --
> > E   import numpy
> > E   from onnx import TensorProto
> > E   from onnx.helper import make_tensor
> > E   from onnxscript import script, external_tensor
> > E   from onnxscript.values import Opset
> > E   from onnxscript.onnx_types import FLOAT
> > E   from onnxscript.onnx_opset import opset13
> > E   
> > E   @script()
> > E   def bck_test_softmax_example(x: FLOAT[1,3]) -> (FLOAT[1,3]):
> > E       y = opset13.Softmax(x)
> > E       return y
> > ```
> 
> </details>

</details>

To view more test analytics, go to the [Test Analytics Dashboard](https://app.codecov.io/gh/microsoft/onnxscript/tests/rama%2Fdc-var)
<sub>📋 Got 3 mins? [Take this short survey](https://forms.gle/BpocVj23nhr2Y45G7) to help us improve Test Analytics.</sub>
author_association: NONE
user: login: gramalingam
created_at: 2025-04-10T16:09:49Z
body: > Could this be extended for _allow_other_attributes as well?

Please clarify: what are you proposing? Thanks!
author_association: COLLABORATOR