[FRONTEND][TFLITE][BugFix] Fix variable typo in batchmatmul converting func (apache#15259)

* TFLite frontend bug fix

* Update tflite.py

* lint

* Add pytest

Author: rebel-shshin

python/tvm/relay/frontend/tflite.py

@@ -3006,7 +3006,9 @@ def convert_batch_matmul(self, op):
                 rank_diff = rank_a - rank_b
                 new_b_shape = _op.concatenate(
                     [
-                        _expr.const([1] * rank_diff, dtype=_infer_type(b_shape).checked_type.dtype),
+                        _expr.const(
+                            [1] * rank_diff, dtype=_infer_type(new_b_shape).checked_type.dtype
+                        ),
                         shape_b,
                     ],
                     0,
@@ -3015,7 +3017,9 @@ def convert_batch_matmul(self, op):
                 rank_diff = rank_b - rank_a
                 new_a_shape = _op.concatenate(
                     [
-                        _expr.const([1] * rank_diff, dtype=_infer_type(a_shape).checked_type.dtype),
+                        _expr.const(
+                            [1] * rank_diff, dtype=_infer_type(new_a_shape).checked_type.dtype
+                        ),
                         shape_a,
                     ],
                     0,
@@ -3041,9 +3045,9 @@ def convert_batch_matmul(self, op):
                 _op.concatenate([out_batch, _op.strided_slice(shape_b, [rank_b - 2], [rank_b])], 0)
             )
             if not tvm.ir.structural_equal(shape_a, a_broadcasted_shape):
-                input_a = _op.transform.broadcast_to(a, a_broadcasted_shape)
+                input_a = _op.transform.broadcast_to(input_a, a_broadcasted_shape)
             if not tvm.ir.structural_equal(shape_b, b_broadcasted_shape):
-                input_b = _op.transform.broadcast_to(b, b_broadcasted_shape)
+                input_b = _op.transform.broadcast_to(input_b, b_broadcasted_shape)
 
             input_a = self.flatten_to_nd(input_a, shape_a, 3)
             input_b = self.flatten_to_nd(input_b, shape_b, 3)

tests/python/frontend/tflite/test_forward.py

@@ -794,6 +794,15 @@ def test_forward_batch_matmul(config):
         adjoint_b=False,
         quantized=config[2],
     )
+    _test_batch_matmul(
+        (2, 3, 5, 4),
+        (1, 3, 5, 4),
+        dtype=config[0],
+        out_dtype=config[1],
+        adjoint_a=True,
+        adjoint_b=False,
+        quantized=config[2],
+    )
     _test_batch_matmul(
         (3, 5, 4),
         (3, 5, 4),
@@ -803,6 +812,15 @@ def test_forward_batch_matmul(config):
         adjoint_b=True,
         quantized=config[2],
     )
+    _test_batch_matmul(
+        (2, 3, 5, 4),
+        (1, 3, 5, 4),
+        dtype=config[0],
+        out_dtype=config[1],
+        adjoint_a=False,
+        adjoint_b=True,
+        quantized=config[2],
+    )
     _test_batch_matmul(
         (3, 4, 5, 6), (3, 4, 6, 5), dtype=config[0], out_dtype=config[1], quantized=config[2]
     )