Implement reduce all/any for non-const axes (#1657)

Signed-off-by: Tom Wildenhain <[email protected]>

Author: TomWildenhain-Microsoft

tests/test_backend.py

@@ -3582,29 +3582,29 @@ def func(x):
         self._run_test_case(func, [_OUTPUT], {_INPUT: input_val})
 
     def test_reduce_all(self):
-        input_val = np.random.randint(0, 2, (10, 20)).astype(np.bool)
+        input_val = np.random.randint(0, 2, (2, 20)).astype(np.bool)
         def func(x):
             res = tf.reduce_all(input_tensor=x, keepdims=False)
             res1 = tf.reduce_all(input_tensor=x, axis=[0], keepdims=False)
             return tf.identity(res, name=_TFOUTPUT), tf.identity(res1, name=_TFOUTPUT1)
         self._run_test_case(func, [_OUTPUT, _OUTPUT1], {_INPUT: input_val})
 
-        input_val = np.random.randint(0, 2, (10, 20)).astype(np.bool)
+        input_val = np.random.randint(0, 2, (2, 20)).astype(np.bool)
         def func(input_x):
             res = tf.reduce_all(input_tensor=input_x, keepdims=True)
             res1 = tf.reduce_all(input_tensor=input_x, axis=[0], keepdims=True)
             return tf.identity(res, name=_TFOUTPUT), tf.identity(res1, name=_TFOUTPUT1)
         self._run_test_case(func, [_OUTPUT, _OUTPUT1], {_INPUT: input_val})
 
     def test_reduce_any(self):
-        input_val = np.random.randint(0, 2, (10, 20)).astype(np.bool)
+        input_val = np.random.randint(0, 2, (2, 20)).astype(np.bool)
         def func(x):
             res = tf.reduce_any(input_tensor=x, keepdims=False)
             res1 = tf.reduce_any(input_tensor=x, axis=[0], keepdims=False)
             return tf.identity(res, name=_TFOUTPUT), tf.identity(res1, name=_TFOUTPUT1)
         self._run_test_case(func, [_OUTPUT, _OUTPUT1], {_INPUT: input_val})
 
-        input_val = np.random.randint(0, 2, (10, 20)).astype(np.bool)
+        input_val = np.random.randint(0, 2, (2, 20)).astype(np.bool)
         def func(x):
             res = tf.reduce_any(input_tensor=x, keepdims=True)
             res1 = tf.reduce_any(input_tensor=x, axis=[0], keepdims=True)
@@ -3613,14 +3613,14 @@ def func(x):
 
     @check_opset_min_version(11, "ReduceMin")
     def test_reduce_all_negative_axis(self):
-        input_val = np.random.randint(0, 2, (10, 20)).astype(np.bool)
+        input_val = np.random.randint(0, 2, (2, 20)).astype(np.bool)
         def func(x):
             res = tf.reduce_all(input_tensor=x, keepdims=False)
             res1 = tf.reduce_all(input_tensor=x, axis=[-1], keepdims=False)
             return tf.identity(res, name=_TFOUTPUT), tf.identity(res1, name=_TFOUTPUT1)
         self._run_test_case(func, [_OUTPUT, _OUTPUT1], {_INPUT: input_val})
 
-        input_val = np.random.randint(0, 2, (10, 20)).astype(np.bool)
+        input_val = np.random.randint(0, 2, (2, 20)).astype(np.bool)
         def func(input_x):
             res = tf.reduce_all(input_tensor=input_x, keepdims=True)
             res1 = tf.reduce_all(input_tensor=input_x, axis=[-1], keepdims=True)
@@ -3629,14 +3629,14 @@ def func(input_x):
 
     @check_opset_min_version(11, "ReduceSum")
     def test_reduce_any_negative_axis(self):
-        input_val = np.random.randint(0, 2, (10, 20)).astype(np.bool)
+        input_val = np.random.randint(0, 2, (2, 20)).astype(np.bool)
         def func(x):
             res = tf.reduce_any(input_tensor=x, keepdims=False)
             res1 = tf.reduce_any(input_tensor=x, axis=[-1], keepdims=False)
             return tf.identity(res, name=_TFOUTPUT), tf.identity(res1, name=_TFOUTPUT1)
         self._run_test_case(func, [_OUTPUT, _OUTPUT1], {_INPUT: input_val})
 
-        input_val = np.random.randint(0, 2, (10, 20)).astype(np.bool)
+        input_val = np.random.randint(0, 2, (2, 20)).astype(np.bool)
         def func(x):
             res = tf.reduce_any(input_tensor=x, keepdims=True)
             res1 = tf.reduce_any(input_tensor=x, axis=[-1], keepdims=True)
@@ -3646,13 +3646,31 @@ def func(x):
     @check_opset_min_version(11, "ReduceSum")
     @check_tf_min_version("1.15")
     def test_reduce_any_empty_axis(self):
-        input_val = np.random.randint(0, 2, (10, 20)).astype(np.bool)
+        input_val = np.random.randint(0, 2, (2, 20)).astype(np.bool)
         def func(x):
             res = tf.reduce_any(input_tensor=x, keepdims=False)
             res1 = tf.reduce_any(input_tensor=x, axis=[], keepdims=False)
             return tf.identity(res, name=_TFOUTPUT), tf.identity(res1, name=_TFOUTPUT1)
         self._run_test_case(func, [_OUTPUT, _OUTPUT1], {_INPUT: input_val})
 
+    def test_reduce_all_scalar_axis(self):
+        input_val = np.random.randint(0, 2, (2, 20)).astype(np.bool)
+        def func(x):
+            res = tf.reduce_all(input_tensor=x, keepdims=False)
+            res1 = tf.reduce_all(input_tensor=x, axis=0, keepdims=False)
+            return tf.identity(res, name=_TFOUTPUT), tf.identity(res1, name=_TFOUTPUT1)
+        self._run_test_case(func, [_OUTPUT, _OUTPUT1], {_INPUT: input_val})
+
+    @check_opset_min_version(13, "ReduceSum")
+    @check_tf_min_version("1.15")
+    def test_reduce_any_nonconst_axis(self):
+        input_val = np.random.randint(0, 2, (2, 20)).astype(np.bool)
+        y_val = np.array([1], np.int32)
+        def func(x, y):
+            res = tf.reduce_any(input_tensor=x, axis=y, keepdims=False)
+            return tf.identity(res, name=_TFOUTPUT)
+        self._run_test_case(func, [_OUTPUT], {_INPUT: input_val, _INPUT1: y_val})
+
     @check_opset_min_version(7, "fill")
     def test_zeros_like(self):
         input_x = np.random.random_sample([10, 20]).astype(np.float32)

tf2onnx/onnx_opset/reduction.py

@@ -144,6 +144,29 @@ def version_6(cls, ctx, node, **kwargs):
         ctx.make_node(op_type="Greater", inputs=[reduce_node_output, zero_node.output[0]],
                       name=node.name, outputs=node.output, shapes=shapes, dtypes=dtypes)
 
+    @classmethod
+    def version_13(cls, ctx, node, **kwargs):
+        keepdims = node.get_attr_value('keep_dims')
+        reduce_input = node.input[0]
+        if node.type == "All":
+            reduce_input = ctx.make_node("Not", [reduce_input]).output[0]
+        cast = ctx.make_node("Cast", inputs=[reduce_input], attr={"to": onnx_pb.TensorProto.FLOAT}).output[0]
+        axes_cast = node.input[1]
+        if ctx.get_rank(axes_cast) == 0:
+            # Unsqueeze scalar axes
+            axes_cast = GraphBuilder(ctx).make_unsqueeze({'data': axes_cast, 'axes': [0]})
+        if ctx.get_dtype(axes_cast) != onnx_pb.TensorProto.INT64:
+            axes_cast = ctx.make_node("Cast", inputs=[axes_cast], attr={"to": onnx_pb.TensorProto.INT64}).output[0]
+        reduce_node_output = GraphBuilder(ctx).make_reduce_sum(
+            {"data": cast, "axes": axes_cast, "keepdims": keepdims, "noop_with_empty_axes": 1},
+            shapes=node.output_shapes, op_name_scope=node.name)
+        zero_node = ctx.make_const(utils.make_name("zero_reduce"), np.array(0, dtype=np.float32))
+        greater_node = ctx.make_node(op_type="Greater", inputs=[reduce_node_output, zero_node.output[0]])
+        result = greater_node.output[0]
+        if node.type == "All":
+            result = ctx.make_node("Not", [result]).output[0]
+        ctx.replace_all_inputs(node.output[0], result)
+
 
 @tf_op("AddN")
 class AddN():