[GitHub] [tvm] ymwangg commented on a change in pull request #8250: [TOPI][batch_matmul] Allow cblas batch_matmul implicit batch_size broadcast

[GitBox <gi...@apache.org>]

ymwangg commented on a change in pull request #8250:
URL: https://github.com/apache/tvm/pull/8250#discussion_r650304289



##########
File path: python/tvm/topi/x86/batch_matmul.py
##########
@@ -162,10 +162,10 @@ def batch_matmul_blas_common(cfg, x, y, out_shape, lib):
     assert len(x.shape) == 3 and len(y.shape) == 3, "only support 3-dim batch_matmul"
     XB, M, XK = get_const_tuple(x.shape)
     YB, N, YK = get_const_tuple(y.shape)
-    assert XB == YB, "batch dimension doesn't match"
+    assert (XB == YB) or (YB == 1) or (XB == 1), "batch dimension doesn't match"

Review comment:
       1. Added this case. 2 Updated.

##########
File path: tests/python/contrib/test_cblas.py
##########
@@ -158,10 +158,11 @@ def test_quantized_matmul_add():
 
 
 def verify_batch_matmul(
-    batch, m, l, n, lib, transa=False, transb=False, iterative=False, dtype="float32"
+    batch_a, batch_b, m, l, n, lib, transa=False, transb=False, iterative=False, dtype="float32"
 ):
-    ashape = (batch, l, n) if transa else (batch, n, l)
-    bshape = (batch, m, l) if transb else (batch, l, m)
+    batch = batch_a

Review comment:
       Fixed. Also changed the extern shape calculation.




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org