comaniac commented on a change in pull request #8250:
URL: https://github.com/apache/tvm/pull/8250#discussion_r650295577
##########
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.From this statement it seems the case of `A(1, M, k) x B(8, N, K)` is
supported. If so, please add a case to the test.
2. Please update the docstring. The shapes of x and y are no longer just
`[batch, ...]`.
##########
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:
If `batch_a == 1` and `batch_b == B`, then `batch` should be `batch_b`?
--
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:
[email protected]
