Make Autograph support in-place assignments with static ``slice(start, stop, step)``

[tzunghanjuang]

Context: Make set_item primitive detect slice. If yes, use dynamic_update_slice. Must be used with https://github.com/PennyLaneAI/diastatic-malt/pull/6.

Update (2024/07/12): We no longer consider dynamic_update_slice as an option since it produces unwanted results if the indices are out-of-range.

Description of the Change:

• Malt package has a new implementation for slice(start, stop, step). We have to capture such arguments.

Possible Drawbacks: If the length of assignment does not match, jax.lax.dynamic_update_slice will adjust the starting point.

• Only static slices are accepted.

result = jnp.empty((6, ), dtype=x.dtype)
result[::2] = jnp.ones((3,))
# result will be [1, 0, 1, 0, 1,0 ]

Related GitHub Issues:

516

[sc-60315]

[dime10]

We might want to push a new release of malt before merging this PR :) @josh146

[dime10]

@tzunghanjuang This needs an update because we decided against using dynamic_update_slice right?

[tzunghanjuang]

@dime10 Yes, I am still looking for the alternative.

[dime10]

@dime10 Yes, I am still looking for the alternative.

Only allowing static indices right?

[tzunghanjuang]

Only allowing static indices right?

Yes

[tzunghanjuang]

@dime10 I have updated the set_item primitive and tests.

[dime10]

The malt package has been updated btw :) https://pypi.org/project/diastatic-malt/

[tzunghanjuang]

HI @dime10. What should be the category for PR? (Feature, improvement or bug fix?)

[dime10]

HI @dime10. What should be the category for PR? (Feature, improvement or bug fix?)

Usually we don't triage PRs by what kind they are (although if you really want to you can). Instead, we assign which part of the project the PR is modifying (docs, ci, frontend, compiler, runtime).

The classification (improvement, bug, etc) is usually used for issues.

EDIT: Oh are you referring to the changelog? In that case I think it is an improvement (since it improves the existing feature array assignment feature) :)

[codecov[bot]]

Codecov Report

All modified and coverable lines are covered by tests :white_check_mark:

Project coverage is 97.99%. Comparing base (93c05e5) to head (8e48834).

Additional details and impacted files

```diff @@ Coverage Diff @@ ## main #843 +/- ## ======================================= Coverage 97.98% 97.99% ======================================= Files 71 71 Lines 10546 10548 +2 Branches 960 961 +1 ======================================= + Hits 10334 10336 +2 Misses 169 169 Partials 43 43 ```

:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.