Improve support for Torch and Jax with `dynamic_one_shot` #5672

mudit2812 · 2024-05-08T20:07:36Z

Context:
Opened in favour of #5630. Bug fix for #5442. This PR updates dynamic_one_shot so that it has better compatibility with the torch and jax interfaces.

Description of the Change:

Change casting method from array.astype() to qml.math.cast in the apply_operation dispatch for MidMeasureMP.
Update usage of qml.math in dynamic_one_shot.
When using qml.counts, cast results to ints before converting to strings for lists of MCM values and floats for single MCM values. This is needed because jax arrays are not hashable, and the hash of torch tensors seems to be independent of the value(s) stored inside it. Thus, neither can be used as keys for dictionaries.

Benefits:
Better interface support with dynamic_one_shot.

Possible Drawbacks:

Related GitHub Issues:

…ples

Co-authored-by: Christina Lee <[email protected]>

…ic_samples_jit

mudit2812 · 2024-05-14T18:36:46Z

[sc-59938]

github-actions · 2024-05-15T20:45:28Z

Hello. You may have forgotten to update the changelog!
Please edit doc/releases/changelog-dev.md 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.

vincentmr

Thanks @mudit2812 . Just a couple suggestions and a kudo ;) I'm ready to approved afterwards.

pennylane/devices/qubit/apply_operation.py

pennylane/transforms/dynamic_one_shot.py

tests/devices/default_qubit/test_default_qubit_native_mcm.py

vincentmr

Looks good to me, thanks @mudit2812 .

dwierichs

Thanks @mudit2812 , this is a neat extension of MCM support! :)

Had three small suggestions, a question for understanding and would like to iterate Vincent's kudo!

doc/releases/changelog-dev.md

pennylane/transforms/dynamic_one_shot.py

tests/devices/default_qubit/test_default_qubit_native_mcm.py

mudit2812 · 2024-05-16T21:17:56Z

I will keep iterating with changing the number of shots/seed until the tests pass, but otherwise this should be good to go.

dwierichs

Thanks @mudit2812 !

codecov · 2024-05-24T22:54:36Z

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 99.67%. Comparing base (fbc2a39) to head (7d1915d).

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #5672      +/-   ##
==========================================
- Coverage   99.68%   99.67%   -0.01%     
==========================================
  Files         416      416              
  Lines       39105    38815     -290     
==========================================
- Hits        38981    38690     -291     
- Misses        124      125       +1

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

mudit2812 and others added 30 commits March 7, 2024 12:37

Added rng and prng_key to get_final_state, apply_operation

2d6d336

Use rng in apply_operation args; linting

7560cdf

WIP

11290a7

Fix measure_with_samples' handling of mid_measurements.

1313af8

Remove comments.

412374e

Merge branch 'master' into feature/dynamic_samples

bd32a35

update changelog

19d7f73

Fix legacy node native mcm test.

8aea51e

Fix old device API

3fb1758

Merge branch 'master' into feature/dynamic_samples

57386e6

Fill out mid_measurements in mock device.

b9d8996

Refactor using masks.

5f7a33f

Merge remote-tracking branch 'origin/master' into feature/dynamic_sam…

146baef

…ples

Update rng use; fix docs

c975493

[skip ci] Skip CI

6899b2c

Always compute all results, even if mv = -1.

87f34b2

[skip ci] testing changes to native MCM tests

721c284

Make post-processing jax-ready. WIP

4aabfe5

Update pennylane/devices/qubit/sampling.py

2614d53

Co-authored-by: Christina Lee <[email protected]>

Implement Christina's suggestions.

c707d46

Merge branch 'master' into feature/dynamic_samples

bd800bb

Add rng and jit mid_measure. WARN: regression

f89139a

Merge branch 'master' into feature/dynamic_samples

f9c7b10

WIP

f7d7fd2

Implement apply_mid_meas with norm of a branch @Mudit.

a836a51

Merge remote-tracking branch 'origin/simulate-rng' into feature/dynam…

5b1b732

…ic_samples_jit

Add jax-jit support in apply_cond/mid_meas.

34b513d

Introduce prng

2562b94

WIP

3db8767

Merge branch 'master' into feature/dynamic_samples

ceaa339

Fixed torch and jax interfaces; added tests

cae872e

mudit2812 marked this pull request as ready for review May 15, 2024 20:45

mudit2812 added 2 commits May 15, 2024 17:10

Small default interface fix; changelog entry

8f6a8cd

Merge branch 'master' into dos-interfaces

148159e

mudit2812 requested a review from vincentmr May 15, 2024 21:11

Added diff_method to tests

a54f703

mudit2812 linked an issue May 15, 2024 that may be closed by this pull request

[BUG] Torch interface incompatibility with finite-shot mid-circuit measurements #5442

Closed

1 task

Reverted change to qml math

2ed91dc

vincentmr reviewed May 16, 2024

View reviewed changes

pennylane/devices/qubit/apply_operation.py Outdated Show resolved Hide resolved

pennylane/transforms/dynamic_one_shot.py Outdated Show resolved Hide resolved

tests/devices/default_qubit/test_default_qubit_native_mcm.py Outdated Show resolved Hide resolved

mudit2812 added 2 commits May 16, 2024 16:23

Fixed tests

75c5031

Merge branch 'master' into dos-interfaces

e19bfc6

mudit2812 requested a review from vincentmr May 16, 2024 20:25

mudit2812 commented May 16, 2024

View reviewed changes

tests/devices/default_qubit/test_default_qubit_native_mcm.py Show resolved Hide resolved

vincentmr approved these changes May 16, 2024

View reviewed changes

dwierichs requested changes May 16, 2024

View reviewed changes

Addressed review

af760f8

mudit2812 requested a review from dwierichs May 16, 2024 21:17

dwierichs approved these changes May 16, 2024

View reviewed changes

mudit2812 changed the title ~~Improve support for interfaces with dynamic_one_shot~~ Improve support for Torch and Jax with dynamic_one_shot May 17, 2024

mudit2812 added 4 commits May 17, 2024 13:45

Fixed torch tests

9cb6d46

Added jax dynamic_one_shot unit tests

ce49af4

Tidy up tests

b45cbcb

Merge branch 'master' into dos-interfaces

7d1915d

mudit2812 enabled auto-merge (squash) May 24, 2024 22:30

mudit2812 merged commit 59a1e05 into master May 24, 2024
40 checks passed

mudit2812 deleted the dos-interfaces branch May 24, 2024 22:54

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Improve support for Torch and Jax with `dynamic_one_shot` #5672

Improve support for Torch and Jax with `dynamic_one_shot` #5672

mudit2812 commented May 8, 2024 •

edited

mudit2812 commented May 14, 2024

github-actions bot commented May 15, 2024

vincentmr left a comment

vincentmr left a comment

dwierichs left a comment

mudit2812 commented May 16, 2024

dwierichs left a comment

codecov bot commented May 24, 2024

Improve support for Torch and Jax with dynamic_one_shot #5672

Improve support for Torch and Jax with dynamic_one_shot #5672

Conversation

mudit2812 commented May 8, 2024 • edited

mudit2812 commented May 14, 2024

github-actions bot commented May 15, 2024

vincentmr left a comment

Choose a reason for hiding this comment

vincentmr left a comment

Choose a reason for hiding this comment

dwierichs left a comment

Choose a reason for hiding this comment

mudit2812 commented May 16, 2024

dwierichs left a comment

Choose a reason for hiding this comment

codecov bot commented May 24, 2024

Codecov Report

Improve support for Torch and Jax with `dynamic_one_shot` #5672

Improve support for Torch and Jax with `dynamic_one_shot` #5672

mudit2812 commented May 8, 2024 •

edited