Autograph automatic conversion of in-place operator-based array updates #1143
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
mehrdad2m
force-pushed
the
ag-op-support-finalize
branch
from
08ee23f
to
d947d4f
Compare
Codecov ReportAll modified and coverable lines are covered by tests ✅
Additional details and impacted files@@ Coverage Diff @@
## main #1143 +/- ##
==========================================
+ Coverage 97.87% 97.89% +0.01%
==========================================
Files 75 76 +1
Lines 10801 10858 +57
Branches 1268 1281 +13
==========================================
+ Hits 10572 10629 +57
Misses 179 179
Partials 50 50 ☔ View full report in Codecov by Sentry. |
paul0403
reviewed
Sep 19, 2024
1. simplified the ag_primitives to only have a update_item_with_op which dispatches to the correct arithmatic operation. 2. Added support for updating a slice using assignment operators 3. Reflected the changes to the documentation 4. Added a test for slice update 5. Improved existing tests to reduce the calls to qjit. 6. Addressed some minor review comments
paul0403
reviewed
Sep 20, 2024
rauletorresc
approved these changes
Sep 20, 2024
paul0403
reviewed
Sep 20, 2024
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.
Context:
#717 added support for converting in-place array updates (
arr[i] = x) into the equivalent JAX traceable code (arr.at[i].set(x)). This PR extends that support to operator assignment array updates.Description of the Change:
AugAssignast nodes assigning to a single index or a slice subscript to calls toupdate_item_with_opupdate_item_with_opmethod that map to the correspondingjax.numpy.ndarray.atequivalent methods for JAX arrays and the normal Python operator assignment otherwisetransform_astinCatalystTransformerto invoke the new converterBenefits: We can use
arr[i] += xinstead ofarr.at[i].add(x).Possible Drawbacks: It would be cleaner to have the new converter live in the DiastaticMalt project.
Related GitHub Issues: #757
Based on the solution presented in this PR: #769
Note that this PR was originally implemented externally by #769. This PR aims to revisit that PR.
[sc-60318]