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 PennyLaneAI/diastatic-malt#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]

Nice! Just needs a changelog and a minimum version bump on the malt dependency :)

[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) :)

[dime10]

Thanks @tzunghanjuang :)

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 97.99%. Comparing base (93c05e5) to head (8e48834).
Report is 227 commits behind head on main.

Additional details and impacted files

@@           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           

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