Merge pull request #1069 from phager90/fix/gemmRewriter_BiasAdd

Fix/gemm rewriter bias add

Author: TomWildenhain-Microsoft

tests/test_backend.py

@@ -2997,6 +2997,21 @@ def func(a, b, c):
         self._run_test_case(func, [_OUTPUT], {_INPUT: x_val1, _INPUT1: x_val2, _INPUT2: x_val3},
                             graph_validator=lambda g: check_op_count(g, "Gemm", 1))
 
+    # test for gemm pattern4: A*B + C [addbias] - 1D bias!
+    def test_gemm_pattern4(self):
+        max_number = 10
+        m = np.random.randint(max_number)
+        n = np.random.randint(max_number)
+        k = np.random.randint(max_number) # bias add requires 1D tensor
+        x_val1 = np.random.rand(m, n).astype("float32")
+        x_val2 = np.random.rand(n, k).astype("float32")
+        x_val3 = np.random.rand(k).astype("float32")
+        def func(a, b, c):
+            x_ = tf.nn.bias_add(tf.matmul(a, b), c)
+            return tf.identity(x_, name=_TFOUTPUT)
+        self._run_test_case(func, [_OUTPUT], {_INPUT: x_val1, _INPUT1: x_val2, _INPUT2: x_val3},
+                            graph_validator=lambda g: check_op_count(g, "Gemm", 1))
+
     # test for gemm pattern0: alpha*A*B + beta*C
     @check_opset_min_version(12, "Optimizer bug in ORT 1.2")
     def test_gemm_pattern0_fail_broadcast(self):

tf2onnx/rewriter/gemm_rewriter.py

@@ -55,7 +55,14 @@ def rewrite_gemm(g, ops):
             OpTypePattern('*', name='C'),
         ])
 
-    pattern_list = [pattern0, pattern1, pattern2, pattern3]
+    # pattern4: A*B + c
+    pattern4 = \
+        OpTypePattern('BiasAdd', name='add', inputs=[
+            OpTypePattern('MatMul', name='matmul'),
+            OpTypePattern('*', name='C'),
+        ])
+
+    pattern_list = [pattern0, pattern1, pattern2, pattern3, pattern4]
 
     for pattern in pattern_list:
         matcher = GraphMatcher(pattern, allow_reorder=True)

tf2onnx/tfonnx.py

@@ -451,13 +451,13 @@ def compat_handler(ctx, node, **kwargs):
     # pre-processing graph rewrites
     # bi-directional re-writer should be placed after single directional re-writer
     rewriters = [rewrite_constant_fold, rewrite_quantize_and_dequantize, rewrite_transpose, rewrite_flatten,
-                 rewrite_gemm, rewrite_random_uniform, rewrite_random_uniform_fold_const,
+                 rewrite_random_uniform, rewrite_random_uniform_fold_const,
                  rewrite_random_normal, rewrite_dropout, rewrite_eye,
                  rewrite_leakyrelu, rewrite_thresholded_relu, rewrite_conv2d_with_pad,
                  rewrite_single_direction_lstm, rewrite_bi_direction_lstm,
                  rewrite_single_direction_gru, rewrite_bi_direction_gru,
                  rewrite_custom_rnn_cell, rewrite_generic_loop, rewrite_cond,
-                 rewrite_biasadd_with_conv2d,
+                 rewrite_biasadd_with_conv2d, rewrite_gemm
                  ]
 
     if custom_rewriter is not None: