Rewriter: don't fuse Min/Max into Clip when a bound has more dims than the input

onnxscript/rewriter/rules/common/_min_max_to_clip.py

@@ -12,6 +12,7 @@
     - fused_const is the reduction (min or max) over all constant inputs.
     - For Clip fusion:
         * All constant inputs must be scalars.
+        * No constant input may have more dimensions than X, since Clip keeps the shape of X.
         * The effective lower bound is the maximum of all lower-bound constants.
         * The effective upper bound is the minimum of all upper-bound constants.
 
@@ -72,14 +73,15 @@ def rewrite(self, op, x, out1, out2):
     def _is_scalar(self, v: np.ndarray) -> bool:
         return np.isscalar(v) or np.size(v) == 1
 
-    def check(self, context, out1, out2, **_):
+    def check(self, context, x, out1, out2, **_):
         """Condition to check if we need to replace the pattern.
 
         Conditions:
             - The min and max input nodes must not be graph inputs.
             - These inputs (except the first) must be constant values (from Constant nodes or initializers).
             - In the case of Min(Max) and Max(Min) patterns:
                 * All inputs must be scalars (as Clip requires scalars).
+                * No input may have more dimensions than X (Clip keeps the shape of X).
                 For Max(Min) pattern:
                     * The lower bound must be less than or equal to the upper bound.
 
@@ -92,15 +94,25 @@ def check(self, context, out1, out2, **_):
 
         first_node = out1.producer()
         second_node = out2.producer()
+        # Min/Max broadcast X up to the rank of a bound such as float[1, 1], but Clip
+        # keeps the shape of X. An unknown rank only allows scalar bounds.
+        x_rank = len(x.shape) if x.shape is not None else 0
 
         # Ensure all inputs except the first are constants
         for input_ in first_node.inputs[1:] + second_node.inputs[1:]:
             if ir.convenience.get_const_tensor(input_) is None:
                 return check_result.fail(f"{input_.name} is not a constant.")
 
             # If scalars are required (Clip fusion), enforce scalar-ness
-            if self.need_scalars and not self._is_scalar(input_.const_value.numpy()):
-                return check_result.fail(f"{input_.name} is not a scalar.")
+            if self.need_scalars:
+                value = input_.const_value.numpy()
+                if not self._is_scalar(value):
+                    return check_result.fail(f"{input_.name} is not a scalar.")
+                if value.ndim > x_rank:
+                    return check_result.fail(
+                        f"{input_.name} has more dimensions than {x.name}, "
+                        "so the fused Clip would change the output shape."
+                    )
 
         if self.need_scalars and self.check_bounds:
             # For Clip fusion in the case of Max(Min(X, upper_bound), lower_bound): check that lower_bound <= upper_bound

onnxscript/rewriter/rules/common/_min_max_to_clip_test.py

@@ -343,6 +343,43 @@ def test_failure_max_min_to_clip_need_scalars(self):
         self.run_failed_condition_test(base_model, max_min_rule, "is not a scalar")
 
 
+class TestClipFusionBoundRank(_TestMinMaxToClipBase):
+    """Min/Max broadcast X up to the rank of the bounds, while Clip keeps the shape of X."""
+
+    @staticmethod
+    def make_model(rule: RewriteRule, signature: str, bound_type: str) -> ir.Model:
+        if rule is min_max_rule:  # Max(Min(X, ub), lb)
+            first, second = "Min(X, ub)", "Max(x1, lb)"
+        else:  # Min(Max(X, lb), ub)
+            first, second = "Max(X, lb)", "Min(x1, ub)"
+        return ir.from_onnx_text(f"""
+            < ir_version: 10, opset_import: ["" : 20] >
+            test_model {signature}
+            <{bound_type} lb = {{0.0}}, {bound_type} ub = {{6.0}}>
+            {{
+                x1 = {first}
+                Y = {second}
+            }}
+        """)
+
+    @parameterized.expand(
+        [
+            ("min_max_1d", min_max_rule, "(float[N] X) => (float[1, N] Y)", "float[1, 1]"),
+            ("min_max_0d", min_max_rule, "(float X) => (float[1] Y)", "float[1]"),
+            ("max_min_1d", max_min_rule, "(float[N] X) => (float[1, N] Y)", "float[1, 1]"),
+            ("max_min_0d", max_min_rule, "(float X) => (float[1] Y)", "float[1]"),
+        ]
+    )
+    def test_failure_clip_fusion_bound_rank_above_input_rank(self, _, rule, signature, bound):
+        base_model = self.make_model(rule, signature, bound)
+        self.run_failed_condition_test(base_model, rule, "has more dimensions than")
+
+    @parameterized.expand([("min_max", min_max_rule), ("max_min", max_min_rule)])
+    def test_successful_clip_fusion_bound_rank_equal_to_input_rank(self, _, rule):
+        base_model = self.make_model(rule, "(float[N] X) => (float[N] Y)", "float[1]")
+        self.run_test(base_model, expected_op_types=["Clip"])
+
+
 class TestIntegrationMinMaxToClip(_TestMinMaxToClipBase):
     def test_successful_full_chain_fusion(self):
         model = ir.from_onnx_text("""