Update

Lint fix

Author: jcf94

python/tvm/relay/frontend/tensorflow.py

@@ -56,7 +56,7 @@
     # Change this flag to False to directly convert to `nn.batch_matmul`.
     # Note that `nn.batch_matmul` with format other than NT is in experimental, it may have some
     # performance issues.
-    "use_nt_batch_matmul_op": True,
+    "use_nt_batch_matmul": True,
 }
 
 # compatible operators that do NOT require any conversion.
@@ -1219,7 +1219,7 @@ def from_tensorflow(self, graph, layout="NHWC", shape=None, outputs=None):
         return func, self._params
 
 
-def from_tensorflow(graph, layout="NHWC", shape=None, outputs=None, convert_config={}):
+def from_tensorflow(graph, layout="NHWC", shape=None, outputs=None, convert_config=None):
     """Load tensorflow graph which is a python tensorflow graph object into relay.
     The companion parameters will be handled automatically.
 
@@ -1237,13 +1237,13 @@ def from_tensorflow(graph, layout="NHWC", shape=None, outputs=None, convert_conf
     outputs : List of output tensor names (Optional)
         if not specified then the last node is assumed as graph output.
 
-    convert_config : Dict[str, Any]
+    convert_config : Optional[Dict[str, Any]]
         Default config:
-            use_dense_op : bool = True
+            use_dense : bool = True
                 Ture to convert `tf.matmul` to `nn.dense`, else to `nn.matmul`.
                 The `nn.dense` op requires the data tensor to be non-transposed and weight tensor
                 to be transposed, may insert extra `transpose` to the original graph.
-            use_nt_batch_matmul_op : bool = True
+            use_nt_batch_matmul : bool = True
                 True to convert `tf.batch_matmul` to `nn.batch_matmul` strict to NT format
                 (transpose_a=False, transpose_b=True).
 

python/tvm/relay/frontend/tensorflow_ops.py

@@ -1176,7 +1176,7 @@ def _impl(inputs, attr, params, mod):
         adj_x = attr["adj_x"]
         adj_y = attr["adj_y"]
 
-        if TF_DEFAULT_CONFIGS["use_nt_batch_matmul_op"]:
+        if TF_DEFAULT_CONFIGS["use_nt_batch_matmul"]:
             # Strictly convert all batch_matmul to NT format
             input_x = _op.transpose(input_x, axes=[0, 2, 1]) if adj_x else input_x
             input_y = _op.transpose(input_y, axes=[0, 2, 1]) if not adj_y else input_y

python/tvm/topi/nn/batch_matmul.py

@@ -16,10 +16,13 @@
 # under the License.
 """Batch matrix multiplication"""
 # pylint: disable=invalid-name
+import logging
 import tvm
 from tvm import te, auto_scheduler
 from ..utils import get_const_tuple
 
+logger = logging.getLogger("topi")
+
 
 def batch_matmul(
     tensor_a,
@@ -94,6 +97,12 @@ def batch_matmul(
         oshape = (batch, XI, YJ)
     if out_dtype is None:
         out_dtype = tensor_a.dtype
+        if tensor_a.dtype != tensor_b.dtype:
+            logger.warning(
+                "tensor_a has different data type with tensor_b: %s, %s",
+                tensor_a.dtype,
+                tensor_b.dtype,
+            )
 
     if (transpose_a, transpose_b) == (True, True):
         compute_lambda = lambda b, i, j: te.sum(

tests/python/frontend/tensorflow/test_forward.py

@@ -124,7 +124,7 @@ def run_tvm_graph(
     disabled_pass=None,
     ignore_in_shape=False,
     serialize=False,
-    convert_config={},
+    convert_config=None,
 ):
     """Generic function to compile on relay and execute on tvm"""
     input_data = convert_to_list(input_data)
@@ -225,7 +225,7 @@ def compare_tf_with_tvm(
     add_shapes_to_graph_def=True,
     targets=None,
     ignore_in_shape=False,
-    convert_config={},
+    convert_config=None,
 ):
     """Generic function to generate and compare tensorflow and TVM output"""
 
@@ -1811,8 +1811,12 @@ def _test_matmul(i, j, k, dtype, outer=None):
 
                 A_np = np.random.uniform(high=5.0, size=A_shape).astype(dtype)
                 B_np = np.random.uniform(high=5.0, size=B_shape).astype(dtype)
-                compare_tf_with_tvm([A_np, B_np], [A.name, B.name], result.name, {"use_dense_op": True})
-                compare_tf_with_tvm([A_np, B_np], [A.name, B.name], result.name, {"use_dense_op": False})
+                compare_tf_with_tvm(
+                    [A_np, B_np], [A.name, B.name], result.name, {"use_dense": True}
+                )
+                compare_tf_with_tvm(
+                    [A_np, B_np], [A.name, B.name], result.name, {"use_dense": False}
+                )
 
 
 def test_forward_matmul():
@@ -1831,10 +1835,10 @@ def _test_batch_matmul(A_shape, B_shape, dtype, adjoint_a=False, adjoint_b=False
         A_np = np.random.uniform(high=5.0, size=A_shape).astype(dtype)
         B_np = np.random.uniform(high=5.0, size=B_shape).astype(dtype)
         compare_tf_with_tvm(
-            [A_np, B_np], [A.name, B.name], result.name, {"use_nt_batch_matmul_op": True}
+            [A_np, B_np], [A.name, B.name], result.name, {"use_nt_batch_matmul": True}
         )
         compare_tf_with_tvm(
-            [A_np, B_np], [A.name, B.name], result.name, {"use_nt_batch_matmul_op": False}
+            [A_np, B_np], [A.name, B.name], result.name, {"use_nt_batch_matmul": False}
         )