[Relax][Op] Add new relax operators #43
Merged
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
* Add LowerToTensorIRPass * Register several high level ops * Register gradients for ops * Unit tests for Ops and gradients [Pass] Add new SimpleAD pass (#3) * Add SimpleAD pass * Write unit tests Co-authored-by: SiriusNEO <1713833595@qq.com> [Relax][AD] Support tuple, modify APIs and move Ops to a new folder (#5) This PR adds support for tuples in the automatic differentiation process of relax. Specifically, two operators, `Tuple()` and `TupleGetItem`, are supported. It also modifies `SimpleAD` api to remove string arguments and use GlobalVar, and adds `gradient` API, a new func to func interface. Besides, we have moved the Ops created by us into a new folder `src/relax/op/training`. [Testing] Relax LSTM demo and numeric gradient test (#6) * lstm test changes * relax training lstm demo Rebase fix (#8) * fix rebase error * rebase fix: AD adapts to new parser * rebase fix: return RuntimeDepShape now Co-authored-by: ubospica <ubospica@gmail.com> [Pass] Polish SimpleAD pass (#9) * SimpleAD now will copy params, which fix the shared Var between functions problem * add error unitests * modify type of require_grads and error test Optimizer API (#7) Co-authored-by: SiriusNEO <1713833595@qq.com> [Relax][Training] Relax Trainer API (#10) * append call pass * add test for append call * successfully train LSTM using this trainer * refractor training structure * add tests for trainer * remove a file * Remove optimizer from this branch * sync with optimizer api * redesign API and add document * fix some problems Co-authored-by: ubospica <ubospica@gmail.com> [FIX][Relax AD] replace ones_like with ones (#11) * ones_like to one * fix rebase Co-authored-by: ubospica <ubospica@gmail.com> Op polish revision
Ubospica
changed the title
- Remove test_relax_ops running part - Add op legalizer test
|
You seemed to have changed the file |
MasterJH5574
approved these changes
Dec 13, 2022
|
@MasterJH5574 Thanks for your review! |
MasterJH5574
pushed a commit
that referenced
this pull request
Dec 14, 2022
* [Op] High level op added * Add LowerToTensorIRPass * Register several high level ops * Register gradients for ops * Unit tests for Ops and gradients [Pass] Add new SimpleAD pass (#3) * Add SimpleAD pass * Write unit tests Co-authored-by: SiriusNEO <1713833595@qq.com> [Relax][AD] Support tuple, modify APIs and move Ops to a new folder (#5) This PR adds support for tuples in the automatic differentiation process of relax. Specifically, two operators, `Tuple()` and `TupleGetItem`, are supported. It also modifies `SimpleAD` api to remove string arguments and use GlobalVar, and adds `gradient` API, a new func to func interface. Besides, we have moved the Ops created by us into a new folder `src/relax/op/training`. [Testing] Relax LSTM demo and numeric gradient test (#6) * lstm test changes * relax training lstm demo Rebase fix (#8) * fix rebase error * rebase fix: AD adapts to new parser * rebase fix: return RuntimeDepShape now Co-authored-by: ubospica <ubospica@gmail.com> [Pass] Polish SimpleAD pass (#9) * SimpleAD now will copy params, which fix the shared Var between functions problem * add error unitests * modify type of require_grads and error test Optimizer API (#7) Co-authored-by: SiriusNEO <1713833595@qq.com> [Relax][Training] Relax Trainer API (#10) * append call pass * add test for append call * successfully train LSTM using this trainer * refractor training structure * add tests for trainer * remove a file * Remove optimizer from this branch * sync with optimizer api * redesign API and add document * fix some problems Co-authored-by: ubospica <ubospica@gmail.com> [FIX][Relax AD] replace ones_like with ones (#11) * ones_like to one * fix rebase Co-authored-by: ubospica <ubospica@gmail.com> Op polish revision * fix ops * debug for ops * prepare for pr * fix pr issues * - Add collapse_sum document - Remove test_relax_ops running part - Add op legalizer test * fix qemu-hack * remove unnecessary imports from test
2 tasks
MasterJH5574
pushed a commit
that referenced
this pull request
Jan 19, 2023
This PR migrates #46 to new struct info infra, as part of our AD migration. Because we need do numerical testing for gradients, this PR depends on the operator legalizer #96. Also because the original version of legalizer did not handle the negative indexing case of `relax.mean`, this PR fixes it. To lower `collapse_sum_to`, `collapse_sum_like` properly, this PR migrates a previous patch #43 which introduces `collapse_sum` in topi. Now we can remove the skip marker in the legalizer test for `collapse_sum_to` and `collapse_sum_like`. The gradients of `cross_entropy` and `softmax_cross_entropy` are removed. And the former will be added back and adjust to new `cross_entropy` introduced in #96. Further plan in this PR: - [x] Add gradients for `log_softmax` and `nll_loss` once #94 is merged. - [x] Gradients for some tuple related operators such as `split` and `concat`. It can help us to test the correctness of AD when there are Tuple-I/O operators. - (Not in this PR) "Undefined Gradient" representation. As we know, the gradients of some operators w.r.t. specified inputs are undefined or meaningless, such as the partial gradient of `indices` in `take(x, indices)`. Relay directly uses `zeros_like` in this case as it won't affect gradient propagation. Another choice is to introduce a dummy Expr named `UndefinedGradient` to represent it. How do we handle this case in relax?
MasterJH5574
pushed a commit
that referenced
this pull request
Jan 28, 2023
This PR migrates #46 to new struct info infra, as part of our AD migration. Because we need do numerical testing for gradients, this PR depends on the operator legalizer #96. Also because the original version of legalizer did not handle the negative indexing case of `relax.mean`, this PR fixes it. To lower `collapse_sum_to`, `collapse_sum_like` properly, this PR migrates a previous patch #43 which introduces `collapse_sum` in topi. Now we can remove the skip marker in the legalizer test for `collapse_sum_to` and `collapse_sum_like`. The gradients of `cross_entropy` and `softmax_cross_entropy` are removed. And the former will be added back and adjust to new `cross_entropy` introduced in #96. Further plan in this PR: - [x] Add gradients for `log_softmax` and `nll_loss` once #94 is merged. - [x] Gradients for some tuple related operators such as `split` and `concat`. It can help us to test the correctness of AD when there are Tuple-I/O operators. - (Not in this PR) "Undefined Gradient" representation. As we know, the gradients of some operators w.r.t. specified inputs are undefined or meaningless, such as the partial gradient of `indices` in `take(x, indices)`. Relay directly uses `zeros_like` in this case as it won't affect gradient propagation. Another choice is to introduce a dummy Expr named `UndefinedGradient` to represent it. How do we handle this case in relax?
MasterJH5574
pushed a commit
that referenced
this pull request
Jan 31, 2023
This PR migrates #46 to new struct info infra, as part of our AD migration. Because we need do numerical testing for gradients, this PR depends on the operator legalizer #96. Also because the original version of legalizer did not handle the negative indexing case of `relax.mean`, this PR fixes it. To lower `collapse_sum_to`, `collapse_sum_like` properly, this PR migrates a previous patch #43 which introduces `collapse_sum` in topi. Now we can remove the skip marker in the legalizer test for `collapse_sum_to` and `collapse_sum_like`. The gradients of `cross_entropy` and `softmax_cross_entropy` are removed. And the former will be added back and adjust to new `cross_entropy` introduced in #96. Further plan in this PR: - [x] Add gradients for `log_softmax` and `nll_loss` once #94 is merged. - [x] Gradients for some tuple related operators such as `split` and `concat`. It can help us to test the correctness of AD when there are Tuple-I/O operators. - (Not in this PR) "Undefined Gradient" representation. As we know, the gradients of some operators w.r.t. specified inputs are undefined or meaningless, such as the partial gradient of `indices` in `take(x, indices)`. Relay directly uses `zeros_like` in this case as it won't affect gradient propagation. Another choice is to introduce a dummy Expr named `UndefinedGradient` to represent it. How do we handle this case in relax?
MasterJH5574
pushed a commit
that referenced
this pull request
Feb 8, 2023
This PR migrates #46 to new struct info infra, as part of our AD migration. Because we need do numerical testing for gradients, this PR depends on the operator legalizer #96. Also because the original version of legalizer did not handle the negative indexing case of `relax.mean`, this PR fixes it. To lower `collapse_sum_to`, `collapse_sum_like` properly, this PR migrates a previous patch #43 which introduces `collapse_sum` in topi. Now we can remove the skip marker in the legalizer test for `collapse_sum_to` and `collapse_sum_like`. The gradients of `cross_entropy` and `softmax_cross_entropy` are removed. And the former will be added back and adjust to new `cross_entropy` introduced in #96. Further plan in this PR: - [x] Add gradients for `log_softmax` and `nll_loss` once #94 is merged. - [x] Gradients for some tuple related operators such as `split` and `concat`. It can help us to test the correctness of AD when there are Tuple-I/O operators. - (Not in this PR) "Undefined Gradient" representation. As we know, the gradients of some operators w.r.t. specified inputs are undefined or meaningless, such as the partial gradient of `indices` in `take(x, indices)`. Relay directly uses `zeros_like` in this case as it won't affect gradient propagation. Another choice is to introduce a dummy Expr named `UndefinedGradient` to represent it. How do we handle this case in relax?
MasterJH5574
pushed a commit
to MasterJH5574/tlc-relax
that referenced
this pull request
Feb 12, 2023
This PR migrates mlc-ai/relax#46 to new struct info infra, as part of our AD migration. Because we need do numerical testing for gradients, this PR depends on the operator legalizer mlc-ai/relax#96. Also because the original version of legalizer did not handle the negative indexing case of `relax.mean`, this PR fixes it. To lower `collapse_sum_to`, `collapse_sum_like` properly, this PR migrates a previous patch mlc-ai/relax#43 which introduces `collapse_sum` in topi. Now we can remove the skip marker in the legalizer test for `collapse_sum_to` and `collapse_sum_like`. The gradients of `cross_entropy` and `softmax_cross_entropy` are removed. And the former will be added back and adjust to new `cross_entropy` introduced in mlc-ai/relax#96. Further plan in this PR: - [x] Add gradients for `log_softmax` and `nll_loss` once mlc-ai/relax#94 is merged. - [x] Gradients for some tuple related operators such as `split` and `concat`. It can help us to test the correctness of AD when there are Tuple-I/O operators. - (Not in this PR) "Undefined Gradient" representation. As we know, the gradients of some operators w.r.t. specified inputs are undefined or meaningless, such as the partial gradient of `indices` in `take(x, indices)`. Relay directly uses `zeros_like` in this case as it won't affect gradient propagation. Another choice is to introduce a dummy Expr named `UndefinedGradient` to represent it. How do we handle this case in relax?
MasterJH5574
pushed a commit
to MasterJH5574/tlc-relax
that referenced
this pull request
Feb 12, 2023
This PR migrates mlc-ai/relax#46 to new struct info infra, as part of our AD migration. Because we need do numerical testing for gradients, this PR depends on the operator legalizer mlc-ai/relax#96. Also because the original version of legalizer did not handle the negative indexing case of `relax.mean`, this PR fixes it. To lower `collapse_sum_to`, `collapse_sum_like` properly, this PR migrates a previous patch mlc-ai/relax#43 which introduces `collapse_sum` in topi. Now we can remove the skip marker in the legalizer test for `collapse_sum_to` and `collapse_sum_like`. The gradients of `cross_entropy` and `softmax_cross_entropy` are removed. And the former will be added back and adjust to new `cross_entropy` introduced in mlc-ai/relax#96. Further plan in this PR: - [x] Add gradients for `log_softmax` and `nll_loss` once mlc-ai/relax#94 is merged. - [x] Gradients for some tuple related operators such as `split` and `concat`. It can help us to test the correctness of AD when there are Tuple-I/O operators. - (Not in this PR) "Undefined Gradient" representation. As we know, the gradients of some operators w.r.t. specified inputs are undefined or meaningless, such as the partial gradient of `indices` in `take(x, indices)`. Relay directly uses `zeros_like` in this case as it won't affect gradient propagation. Another choice is to introduce a dummy Expr named `UndefinedGradient` to represent it. How do we handle this case in relax?
MasterJH5574
pushed a commit
that referenced
this pull request
Feb 12, 2023
This PR migrates #46 to new struct info infra, as part of our AD migration. Because we need do numerical testing for gradients, this PR depends on the operator legalizer #96. Also because the original version of legalizer did not handle the negative indexing case of `relax.mean`, this PR fixes it. To lower `collapse_sum_to`, `collapse_sum_like` properly, this PR migrates a previous patch #43 which introduces `collapse_sum` in topi. Now we can remove the skip marker in the legalizer test for `collapse_sum_to` and `collapse_sum_like`. The gradients of `cross_entropy` and `softmax_cross_entropy` are removed. And the former will be added back and adjust to new `cross_entropy` introduced in #96. Further plan in this PR: - [x] Add gradients for `log_softmax` and `nll_loss` once #94 is merged. - [x] Gradients for some tuple related operators such as `split` and `concat`. It can help us to test the correctness of AD when there are Tuple-I/O operators. - (Not in this PR) "Undefined Gradient" representation. As we know, the gradients of some operators w.r.t. specified inputs are undefined or meaningless, such as the partial gradient of `indices` in `take(x, indices)`. Relay directly uses `zeros_like` in this case as it won't affect gradient propagation. Another choice is to introduce a dummy Expr named `UndefinedGradient` to represent it. How do we handle this case in relax?
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.
This PR adds these operators:
relax.op.tensorrelax.negativerelax.logrelax.tanhrelax.op.transformrelax.full_likerelax.onesrelax.zerosrelax.ones_likerelax.zeros_likerelax.collapse_sum_likerelax.collapse_sum_totopi.reductiontopi.collapse_sumrelax.collapse_sum_likeandrelax.collapse_sum_to.Besides, this pull request fixes a bug in
python/tvm/relax/block_builder.pyso now BlockBuilder could support functions without parameters.