[SymbolicShapeInference] Support 10 Reduce* ops

onnxruntime/python/tools/symbolic_shape_infer.py

@@ -165,8 +165,16 @@
             "Pad": self._infer_Pad,
             "Range": self._infer_Range,
             "Reciprocal": self._pass_on_shape_and_type,
-            "ReduceSum": self._infer_ReduceSum,
-            "ReduceProd": self._infer_ReduceProd,
+            "ReduceL1": self._infer_Reduce,
+            "ReduceL2": self._infer_Reduce,
+            "ReduceLogSum": self._infer_Reduce,
+            "ReduceLogSumExp": self._infer_Reduce,
+            "ReduceMax": self._infer_Reduce,
+            "ReduceMin": self._infer_Reduce,
+            "ReduceMean": self._infer_Reduce,
+            "ReduceProd": self._infer_Reduce,
+            "ReduceSum": self._infer_Reduce,
+            "ReduceSumSquare": self._infer_Reduce,
             "Reshape": self._infer_Reshape,
             "Resize": self._infer_Resize,
             "Round": self._pass_on_shape_and_type,
@@ -1567,46 +1575,58 @@
             )
         )
 
-    def _infer_ReduceSum(self, node):  # noqa: N802
+    # This func takes care of Reduce*** ops, 
+    # including ReduceSum, ReduceMean, ReduceMin, ReduceMax, ReduceProd, etc
+    def _infer_Reduce(self, node):  # noqa: N802
         keep_dims = get_attribute(node, "keepdims", 1)
-        if get_opset(self.out_mp_) >= 13 and len(node.input) > 1:
-            # ReduceSum changes axes to input[1] in opset 13
+        opset = get_opset(self.out_mp_)
+
+        if opset >= 13 and len(node.input) > 1:
             axes = self._try_get_value(node, 1)
-            vi = self.known_vi_[node.output[0]]
-            if axes is None:
-                assert keep_dims  # can only handle keep_dims==True when axes is unknown, by generating new ranks
-                vi.CopyFrom(
-                    helper.make_tensor_value_info(
-                        node.output[0],
-                        self.known_vi_[node.input[0]].type.tensor_type.elem_type,
-                        get_shape_from_sympy_shape(self._new_symbolic_shape(self._get_shape_rank(node, 0), node)),
-                    )
+        else:
+            axes = get_attribute(node, "axes")
+
+        vi = self.known_vi_[node.output[0]]
+
+        if axes is None:
+            assert keep_dims == 1, f"{node.op_type} Op: Cannot infer shape when axes is unknown and keepdims is not 1."
+            rank = self._get_shape_rank(node, 0)
+            new_shape = self._new_symbolic_shape(rank, node)
+            vi.CopyFrom(
+                helper.make_tensor_value_info(
+                    node.output[0],
+                    self.known_vi_[node.input[0]].type.tensor_type.elem_type,
+                    get_shape_from_sympy_shape(new_shape),
                 )
-            else:
-                shape = self._get_shape(node, 0)
-                output_shape = []
-                axes = [handle_negative_axis(a, len(shape)) for a in axes]
-                for i, d in enumerate(shape):
-                    if i in axes:
-                        if keep_dims:
-                            output_shape.append(1)
+            )
+        else:
+            # Special optimization for ReduceProd
+            if node.op_type == "ReduceProd" and keep_dims == 0 and axes == [0]:
+                data = self._get_int_or_float_values(node)[0]
+                if data is not None:
+                    self.sympy_data_[node.output[0]] = sympy_reduce_product(data)
+                    return
+
+            input_shape = self._get_shape(node, 0)
+            assert input_shape, f"{node.op_type} Op: Reduction over an empty set of values yields undefined"
+
+            axes = [handle_negative_axis(a, len(input_shape)) for a in axes]
+            output_shape = []
+            for i, dim in enumerate(input_shape):
+                if i in axes:
+                    if keep_dims == 1:
+                        output_shape.append(1)
                     else:
-                        output_shape.append(d)
-                vi.CopyFrom(
-                    helper.make_tensor_value_info(
-                        node.output[0],
-                        self.known_vi_[node.input[0]].type.tensor_type.elem_type,
-                        output_shape,
-                    )
+                        continue
+                else:
+                    output_shape.append(dim)
+            vi.CopyFrom(
+                helper.make_tensor_value_info(
+                    node.output[0],
+                    self.known_vi_[node.input[0]].type.tensor_type.elem_type,
+                    output_shape,
                 )
-
-    def _infer_ReduceProd(self, node):  # noqa: N802
-        axes = get_attribute(node, "axes")
-        keep_dims = get_attribute(node, "keepdims", 1)
-        if keep_dims == 0 and axes == [0]:
-            data = self._get_int_or_float_values(node)[0]
-            if data is not None:
-                self.sympy_data_[node.output[0]] = sympy_reduce_product(data)
+            )
 
     def _infer_RelativePositionBias(self, node):  # noqa: N802
         seq_len = self._try_get_value(node, 1)