This is an automated email from the ASF dual-hosted git repository.
masahi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 6bcc1b44f5 [TOPI] Fix dynamic dimensions support for Dense on TOPI
side (#15018)
6bcc1b44f5 is described below
commit 6bcc1b44f57d4541aad7f8fddbbdfd1e59cd8a3b
Author: Valery Chernov <[email protected]>
AuthorDate: Sun Jun 4 15:44:39 2023 +0400
[TOPI] Fix dynamic dimensions support for Dense on TOPI side (#15018)
fix comparison for any_dim
Co-authored-by: Valery Chernov <[email protected]>
---
python/tvm/topi/nn/dense.py | 12 +++++++++---
1 file changed, 9 insertions(+), 3 deletions(-)
diff --git a/python/tvm/topi/nn/dense.py b/python/tvm/topi/nn/dense.py
index d7d475fb0c..ce3aebadb6 100644
--- a/python/tvm/topi/nn/dense.py
+++ b/python/tvm/topi/nn/dense.py
@@ -94,9 +94,15 @@ def matmul(
red_dim, out_dim = tensor_b.shape
# cmp should be done by values
- assert int(in_dim) == int(
- red_dim
- ), "Inner dimensions of dense do not match. {in_dim} vs {red_dim}."
+ condition = True
+ if isinstance(in_dim, tvm.tir.SizeVar): # "any_dim"
+ condition = False
+ elif isinstance(red_dim, tvm.tir.SizeVar): # "any_dim"
+ condition = False
+ if condition:
+ assert int(in_dim) == int(
+ red_dim
+ ), "Inner dimensions of dense do not match. {in_dim} vs {red_dim}."
k = te.reduce_axis((0, in_dim), name="k")
if (transpose_a, transpose_b) == (True, True):
