[Relax] Easier use of Function.bind_params with relax.op.create.* #17323
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
The `Function.bind_params` method requires all bound values to have `StructInfo` annotations. While these are generated on construction for `relax::Constant` and `relax::PrimValue`, the most common cases of bound values, other useful expression such as `relax.op.zeros(shape,dtype)` do not have `StructInfo` until they are normalized. In addition, because `Function.bind_params` may be used in contexts that do not have a `relax.BlockBuilder` available, performing this normalization is not straightforward. This commit updates the constructors for Relax operations within the `tvm.relax.op.create` namespace to infer their `StructInfo` on construction, if it is possible to do so. This allows these operations to be used inside `Function.bind_params`.
Lunderberg
force-pushed
the
relax_allow_op_create_expressions_in_function_bind_params
branch
from
35c0d1c
to
b40070d
Compare
|
Rebase onto main, since the CI failures were caused by #17371. |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
The
Function.bind_paramsmethod requires all bound values to haveStructInfoannotations. While these are generated on construction forrelax::Constantandrelax::PrimValue, the most common cases of bound values, other useful expression such asrelax.op.zeros(shape,dtype)do not haveStructInfountil they are normalized. In addition, becauseFunction.bind_paramsmay be used in contexts that do not have arelax.BlockBuilderavailable, performing this normalization is not straightforward.This commit updates the constructors for Relax operations within the
tvm.relax.op.createnamespace to infer theirStructInfoon construction, if it is possible to do so. This allows these operations to be used insideFunction.bind_params.