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.

Sign up for GitHub

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

Jump to bottom

Add support to AG for Jax array single assignment #717

Merged
merged 3 commits into from
May 8, 2024
Merged

Add support to AG for Jax array single assignment #717

merged 3 commits into from
May 8, 2024

Conversation

rauletorresc
Copy link
Contributor

@rauletorresc rauletorresc commented May 6, 2024 •
edited
Loading

Context: We want to support NumPy-style in-place array updates, such as arr[i] = x

Description of the Change: Overload set_item function from Autograph and enable converter.Feature.LISTS option.

Benefits: We can use directly arr[i] = x instead of arr = arr.at[i].set(x)

Related GitHub Issues: #516

Based on the solution presented in this PR: #582

[sc-60313]

@rauletorresc rauletorresc added good first issue Good for newcomers frontend Pull requests that update the frontend labels May 6, 2024
@rauletorresc rauletorresc self-assigned this May 6, 2024
@rauletorresc rauletorresc force-pushed the autograph_single_index_arrays branch 3 times, most recently from 559a59f to 2c22446 Compare May 6, 2024 17:18
@rauletorresc rauletorresc marked this pull request as ready for review May 6, 2024 17:22
dime10
dime10 reviewed May 7, 2024
View reviewed changes
Copy link
Collaborator

@dime10 dime10 left a comment

Choose a reason for hiding this comment

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

Hi Raul, we should also test setting a dynamic index :)

@rauletorresc rauletorresc force-pushed the autograph_single_index_arrays branch 2 times, most recently from 2fd85a8 to 3710821 Compare May 8, 2024 09:02
@rauletorresc
Copy link
Contributor Author

Hi Raul, we should also test setting a dynamic index :)

Done

@rauletorresc rauletorresc requested a review from dime10 May 8, 2024 09:04
@codecov Codecov
Copy link

codecov bot commented May 8, 2024 •
edited
Loading

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 98.08%. Comparing base (c07baae) to head (93d8705).

❗ Current head 93d8705 differs from pull request most recent head cc840ee. Consider uploading reports for the commit cc840ee to get more accurate results

Additional details and impacted files 
@@           Coverage Diff           @@
##             main     #717   +/-   ##
=======================================
  Coverage   98.08%   98.08%           
=======================================
  Files          69       69           
  Lines        9468     9474    +6     
  Branches      746      747    +1     
=======================================
+ Hits         9287     9293    +6     
  Misses        147      147           
  Partials       34       34

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

@rauletorresc rauletorresc force-pushed the autograph_single_index_arrays branch from 3710821 to d85a5b8 Compare May 8, 2024 09:34
dime10
dime10 approved these changes May 8, 2024
View reviewed changes
Copy link
Collaborator

@dime10 dime10 left a comment

Choose a reason for hiding this comment

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

Looks good to me :) Just needs a changelog entry

frontend/test/pytest/test_autograph.py Outdated Show resolved Hide resolved
@rauletorresc
Copy link
Contributor Author

Looks good to me :) Just needs a changelog entry

I was thinking of adding such entry once all in place array assignment cases are covered. What do you think?

@dime10
Copy link
Collaborator

dime10 commented May 8, 2024 •
edited
Loading

Looks good to me :) Just needs a changelog entry

I was thinking of adding such entry once all in place array assignment cases are covered. What do you think?

You can still create one now, since this PR should be attached to the changelog entry anyways.
What other cases were you thinking of adding?

@rauletorresc
Copy link
Contributor Author

Looks good to me :) Just needs a changelog entry

I was thinking of adding such entry once all in place array assignment cases are covered. What do you think?

You can still created now, since this PR should be attached to the changelog entry anyways. What other cases were you thinking of adding?

  • support for slices
  • operator-based array updates

@rauletorresc rauletorresc force-pushed the autograph_single_index_arrays branch from d85a5b8 to 8d08ebf Compare May 8, 2024 20:15
@rauletorresc
Copy link
Contributor Author

Looks good to me :) Just needs a changelog entry

I was thinking of adding such entry once all in place array assignment cases are covered. What do you think?

You can still create one now, since this PR should be attached to the changelog entry anyways. What other cases were you thinking of adding?

Done!

@rauletorresc rauletorresc requested a review from dime10 May 8, 2024 20:18
dime10
dime10 approved these changes May 8, 2024
View reviewed changes
Copy link
Collaborator

@dime10 dime10 left a comment

Choose a reason for hiding this comment

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

Looks good 👍

@rauletorresc rauletorresc force-pushed the autograph_single_index_arrays branch from 8d08ebf to 93d8705 Compare May 8, 2024 20:31
@rauletorresc rauletorresc force-pushed the autograph_single_index_arrays branch from 93d8705 to cc840ee Compare May 8, 2024 20:47
@rauletorresc rauletorresc merged commit 2258e8b into main May 8, 2024
37 checks passed
@rauletorresc rauletorresc deleted the autograph_single_index_arrays branch May 8, 2024 21:11
@Spencer-Comin Spencer-Comin mentioned this pull request May 28, 2024
Closed
@mehrdad2m mehrdad2m mentioned this pull request Sep 19, 2024
Merged
mehrdad2m added a commit that referenced this pull request Sep 20, 2024
@mehrdad2m @Spencer-Comin
Autograph automatic conversion of in-place operator-based array updat…
621c027 
…es (#1143)

**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:**

- Add new Autograph converter to map `AugAssign` ast nodes assigning to
a single index or a slice subscript to calls to `update_item_with_op`
- Implement `update_item_with_op` method that map to the corresponding
`jax.numpy.ndarray.at` equivalent methods for JAX arrays and the normal
Python operator assignment otherwise
- Overload `transform_ast` in `CatalystTransformer` to invoke the new
converter

**Benefits:** We can use `arr[i] += x` instead of `arr.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.

---------

Co-authored-by: Spencer Comin <scomin@me.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@dime10 dime10 dime10 approved these changes

@sergei-mironov sergei-mironov Awaiting requested review from sergei-mironov

@paul0403 paul0403 Awaiting requested review from paul0403

@erick-xanadu erick-xanadu Awaiting requested review from erick-xanadu

Assignees

@rauletorresc rauletorresc

Labels
frontend Pull requests that update the frontend good first issue Good for newcomers
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
2 participants
@rauletorresc @dime10