Skip to content
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

Merged
merged 3 commits into from
Jul 2, 2021

Conversation

lygztq
Copy link
Contributor

@lygztq lygztq commented Jun 30, 2021

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 of batch_matmul's type relation. When the input tensor has a dynamic shape but not all dimensions are dynamic, this relation function will use x's shape directly and output wrong results.

For example, given x with shape (Any(), 16, 32) and y with shape (Any(), 16, 32), the output shape of current implementation is (Any(), 16, 32).

@lygztq lygztq marked this pull request as ready for review June 30, 2021 13:06
@lygztq
Copy link
Contributor Author

lygztq commented Jun 30, 2021

Comment on lines 382 to 383
# TODO(mbrookhart): enable once VM supports heterogenous execution
# @tvm.testing.uses_gpu
Copy link
Contributor

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

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm not sure, I think this line (which is called by topi.cuda here) will fail if x.shape[2] is static and y.shape[2] is dynamic.

By the way, the unit test in ci reported a strange error. I’m not sure if the error is related to this PR.

Copy link
Contributor

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?

Copy link
Contributor Author

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.

@masahi masahi merged commit ab01abc into apache:main Jul 2, 2021
ylc pushed a commit to ylc/tvm that referenced this pull request Sep 29, 2021
* fix type relation for batch_matmul

* fix lint
zxy844288792 pushed a commit to zxy844288792/tvm that referenced this pull request Mar 4, 2022
* fix type relation for batch_matmul

* fix lint
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants