Fix collapse slices rewrite rules to handle unknown dims

Signed-off-by: Justin Chu <[email protected]>

Author: justinchuby

noxfile.py

@@ -42,7 +42,7 @@
     "packaging",
     "protobuf",
 )
-ONNX_IR = "onnx_ir==0.1.9"
+ONNX_IR = "onnx_ir==0.1.10"
 ONNX_IR_MAIN = "git+https://github.com/onnx/ir-py.git@main#egg=onnx_ir"
 
 

onnxscript/rewriter/rules/common/_collapse_slices.py

@@ -85,6 +85,10 @@ def _same_shape(op, data: ir.Value, slice_output: ir.Value, steps: ir.Value, **_
     if not is_singleton_value(steps, 1):
         return False
 
+    # If any dim is unknown, the shapes are not the same
+    if data.shape.has_unknown_dim() or slice_output.shape.has_unknown_dim():
+        return False
+
     return data.shape == slice_output.shape
 
 

pyproject.toml

@@ -27,7 +27,7 @@ classifiers = [
 dependencies = [
   "ml_dtypes",
   "numpy",
-  "onnx_ir>=0.1.9,<2",  # Expect onnx_ir to have a breaking change in 2.0. If not, extend this range.
+  "onnx_ir>=0.1.10,<2",  # Expect onnx_ir to have a breaking change in 2.0. If not, extend this range.
   "onnx>=1.16",
   "packaging",
   "typing_extensions>=4.10",