Make Autograph ``get_item`` handle list with jax index

PennyLaneAI / catalyst

A JIT compiler for hybrid quantum programs in PennyLane

https://docs.pennylane.ai/projects/catalyst

Apache License 2.0

101 stars 26 forks source link

Make Autograph ``get_item`` handle list with jax index #886

Closed tzunghanjuang closed 2 days ago

tzunghanjuang commented 1 week ago

Context: Autograph passes jax index to list or range and raises the following error.

UserWarning: Tracing of an AutoGraph converted for loop failed with an exception: TracerIntegerConversionError:  The __index__() method was called on traced array with shape int64[].

To fix this, we introduces a conversion in get_item in Autograph to convert list or range into a jax numpy array.

Related GitHub Issues: [sc-67797]

github-actions[bot] commented 1 week ago

Hello. You may have forgotten to update the changelog! Please edit doc/changelog.md on your branch with:

A one-to-two sentence description of the change. You may include a small working example for new features.
A link back to this PR.
Your name (or GitHub username) in the contributors section.

codecov[bot] commented 5 days ago

Codecov Report

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

Project coverage is 97.96%. Comparing base (4894806) to head (2037c01). Report is 2 commits behind head on main.

Additional details and impacted files

```diff @@ Coverage Diff @@ ## main #886 +/- ## ========================================== - Coverage 97.97% 97.96% -0.01% ========================================== Files 71 71 Lines 10449 10465 +16 Branches 948 950 +2 ========================================== + Hits 10237 10252 +15 Misses 169 169 - Partials 43 44 +1 ```

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

dime10 commented 2 days ago

The consensus is that this would go far in terms of transforming the user program in unexpected ways. Closing for now, but it was a good idea @tzunghanjuang!