-
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
[TFLite] Support for BATCH_MATMUL tflite operator #14423
Conversation
Adds support for BATCH_MATMUL operator in the TFLite frontend. Adds a test that checks supported TFLite types.
Thanks for contributing to TVM! Please refer to the contributing guidelines https://tvm.apache.org/docs/contribute/ for useful information and tips. Please request code reviews from Reviewers by @-ing them in a comment. Generated by tvm-bot |
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.
LGTM, thanks @neildhickey!
rank_diff = rank_a - rank_b | ||
new_b_shape = _op.concatenate( | ||
[ | ||
_expr.const([1] * rank_diff, dtype=_infer_type(b_shape).checked_type.dtype), |
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.
the "b_shape" here is a mistake? should be "shape_b"?
rank_diff = rank_b - rank_a | ||
new_a_shape = _op.concatenate( | ||
[ | ||
_expr.const([1] * rank_diff, dtype=_infer_type(a_shape).checked_type.dtype), |
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.
the "a_shape" here is a mistake? should be "shape_a"?
) | ||
) | ||
if not tvm.ir.structural_equal(shape_a, a_broadcasted_shape): | ||
input_a = _op.transform.broadcast_to(a, a_broadcasted_shape) |
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.
"a" is not define?
if not tvm.ir.structural_equal(shape_a, a_broadcasted_shape): | ||
input_a = _op.transform.broadcast_to(a, a_broadcasted_shape) | ||
if not tvm.ir.structural_equal(shape_b, b_broadcasted_shape): | ||
input_b = _op.transform.broadcast_to(b, b_broadcasted_shape) |
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.
"b" is not define
Adds support for BATCH_MATMUL operator in the TFLite frontend.
Adds a test that checks supported TFLite types.