-
Notifications
You must be signed in to change notification settings - Fork 3.5k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
[BugFix][Relay] Fix type relation for batch_matmul #8376
Conversation
# TODO(mbrookhart): enable once VM supports heterogenous execution | ||
# @tvm.testing.uses_gpu |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I thought I'd removed all of these TODO's. Oops! This should work on GPU now, I think
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This is a flaky test that times out sometimes, rebase?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Everything is fine now, thanks.
* fix type relation for batch_matmul * fix lint
* fix type relation for batch_matmul * fix lint
I find the current implementation of
relay.nn.batch_matmul
will fail given input tensors with shape(Any(), M, K)
and(Any(), N, K)
. This is caused by the wrong implementation ofbatch_matmul
's type relation. When the input tensor has a dynamic shape but not all dimensions are dynamic, this relation function will usex
's shape directly and output wrong results.For example, given
x
with shape(Any(), 16, 32)
andy
with shape(Any(), 16, 32)
, the output shape of current implementation is(Any(), 16, 32)
.