Speed up dispatching with autoray #4261

frederikwilde · 2023-06-16T16:02:23Z

Context:
qml.math uses autoray to operate interface independently. When the number of calls to functions in qml.math is large, this can create a significant performance overhead, even if the overhead in each function call is miniscule. The main function being called in autoray is the dispatcher do().

Description of the Change:
By creating a profile for various device, interface, diff-method combinations the following overheads were identified and removed:

Several functions are custom registered and in the case of autograd were importing in place. The in-place import now uses caching.
ndim and shape were previously custom registered, but are now directly imported from autoray.

Benefits:
Faster execution on all devices, particularly for diff_method="parameter-shift".

Comparing to a recent version of the master branch we see a significant speed up for most scenarios. Benchmarks for computing the gradient of two qml.StronglyEntanglingLayers on 10 qubits. w/ caching refers to the case when we pass cache=None to the QNode, which reduces the number of operator hashes (which call autoray.do) significantly!

For details see the Notion doc.

Possible Drawbacks:
More hard-coded handling of qml.math calls.

Related GitHub Issues:
n/a

…lace import module import.

frederikwilde · 2023-06-16T16:03:02Z

[sc-38644]

github-actions · 2023-06-16T16:03:20Z

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.

…rayBox operations. Disable _check_batching for set_parameters in copied operators to reduce do overheads.

* pylint errors * update changelog * tidy up * hopefully fix formatting complaints * hopefully fix formatting complaints * Update doc/releases/changelog-dev.md * formatting and codecov * leave original jordan_wigner docstring in place so qchem documentation is unchanged * Update pennylane/fermi/conversion.py * fix docstring issue with dispatch * Revert "fix docstring issue with dispatch" This reverts commit 9f3e736.

* Updated qcut * Updated qcut tests * Updated changelog * Makde changelog entry moree descriptive * Add the `FermiSentence` class (#4195) * create fermiword class * create fermisentence class * [skip ci] modify order * [skip ci] remove unused import * [skip ci] add deepcopy method * [skip ci] replace list with generator * [skip ci] add tests * [skip ci] add tests * [skip ci] add tests * add tests * [skip ci] run black * [skip ci] copy fermiword * [skip ci] rebase and fix conflict * [skip ci] run black * [skip ci] copy fermiword * [skip ci] add tests * add test to tests_passing_pylint * add sum mult test * [skip ci] add tests * [skip ci] add pow tests * add tests * fix pylint * add error tests * fix pylint * add TODO * modify docstrings * update changelog * update changelog * modify docstrings * fix codecov * add code review comments * add test for indices * fix codefactor * Update pennylane/fermi/fermionic.py Co-authored-by: Jay Soni <jbsoni@uwaterloo.ca> * black * update todos * update tests * Apply suggestions from code review * handle grouping in product of FermiSentences * Update pennylane/fermi/fermionic.py Co-authored-by: Jay Soni <jbsoni@uwaterloo.ca> * refactor __mul__ * Apply suggestions from code review Co-authored-by: Jay Soni <jbsoni@uwaterloo.ca> --------- Co-authored-by: lillian542 <38584660+lillian542@users.noreply.github.com> Co-authored-by: Jay Soni <jbsoni@uwaterloo.ca> Co-authored-by: lillian542 <Lillian.frederiksen@xanadu.ai> * Update Docker Plugins (#4178) * Removed sf and added quantuminspire plugins * Updated interface versions to those in CI. * Updatede interface-gpu versions * Updated Jaxlib for GPU to 0.4.10 * Update docker/interfaces/install-interface-gpu.sh Co-authored-by: Matthew Silverman <ma.silv11@gmail.com> --------- Co-authored-by: Matthew Silverman <ma.silv11@gmail.com> * Fix adjoint jacobian with grad on execution false (#4217) * Fix * Update * Add changelog * Update pennylane/interfaces/execution.py Co-authored-by: David Wierichs <david.wierichs@xanadu.ai> * Update doc/releases/changelog-dev.md Co-authored-by: David Wierichs <david.wierichs@xanadu.ai> * Black --------- Co-authored-by: David Wierichs <david.wierichs@xanadu.ai> * Trigger CI * Update doc/releases/changelog-dev.md * Updating to fix docs build * Further changes to fix doc build * Fixed KaHyPar tests * Update pennylane/operation.py * Added legacy tests * Added new eq and hash to test changes * Testing to see if anything fails with new hash * Reverted changes to op hash * Testing what updated equality breaks * Reverted op.__eq__ * Update docstrings to reflect changes to nodes * Update pennylane/transforms/qcut/processing.py Co-authored-by: Matthew Silverman <matthews@xanadu.ai> --------- Co-authored-by: soranjh <40344468+soranjh@users.noreply.github.com> Co-authored-by: lillian542 <38584660+lillian542@users.noreply.github.com> Co-authored-by: Jay Soni <jbsoni@uwaterloo.ca> Co-authored-by: lillian542 <Lillian.frederiksen@xanadu.ai> Co-authored-by: Manul Patel <77568048+manulpatel@users.noreply.github.com> Co-authored-by: Matthew Silverman <ma.silv11@gmail.com> Co-authored-by: Romain Moyard <rmoyard@gmail.com> Co-authored-by: David Wierichs <david.wierichs@xanadu.ai> Co-authored-by: Matthew Silverman <matthews@xanadu.ai>

* create fermiword class * create fermisentence class * [skip ci] modify order * [skip ci] remove unused import * [skip ci] add deepcopy method * [skip ci] replace list with generator * [skip ci] add tests * [skip ci] add tests * [skip ci] add tests * add tests * [skip ci] run black * [skip ci] copy fermiword * [skip ci] rebase and fix conflict * [skip ci] run black * [skip ci] copy fermiword * [skip ci] add tests * add test to tests_passing_pylint * add sum mult test * [skip ci] add tests * [skip ci] add pow tests * add tests * fix pylint * add error tests * fix pylint * add TODO * modify docstrings * update changelog * update changelog * modify docstrings * fix codecov * add code review comments * add test for indices * fix codefactor * multiply FermiWord and FermiSentence * add tests for fw * fs * add and test multiplying fermiword with integer or float * allow number times fermi sentence * update tests * Allow subtracting one FermiSentence from another * Allow substraction * clean up formatting and tests * Add multiplication by complex * Add more tests and reorganize * update change log * add test * black formatting * Add constants to FermiWord and FermiSentence * add tests * radd and rsub methods plus tests * missing lines for codecov * start adding support for tensors * numpy and pennylane numpy tensors plus tests for __add__ and __sub__ * add tests for multiplying by tensor * fix formatting issues post merge * add comment regarding overriding numpy dunders * Apply suggestions from code review Co-authored-by: soranjh <40344468+soranjh@users.noreply.github.com> * check len of array before adding to FeriSentence * Raise error if array with len>2 is passed to arithmetic dunders * Apply suggestions from code review Co-authored-by: Jay Soni <jbsoni@uwaterloo.ca> * Update doc/releases/changelog-dev.md * round to 10 digits --------- Co-authored-by: soranjh <soranjh@yahoo.com> Co-authored-by: soranjh <40344468+soranjh@users.noreply.github.com> Co-authored-by: Jay Soni <jbsoni@uwaterloo.ca>

* Support broadcasting in state measurements * docs for is_state_batched * Support broadcasting in sample measurements * Apply suggestions from code review Co-authored-by: Christina Lee <christina@xanadu.ai> * black * black * Remove total_copies * fix --------- Co-authored-by: Christina Lee <christina@xanadu.ai>

…irectly (#4241) * introduce warning and tests * changelog, recommendation * docstrings * fix test * switch to raising an error instead * Apply suggestions from code review Co-authored-by: Tom Bromley <49409390+trbromley@users.noreply.github.com> * fix tests --------- Co-authored-by: Korbinian Kottmann <Korbinian.Kottmann@gmail.com> Co-authored-by: Tom Bromley <49409390+trbromley@users.noreply.github.com>

* copy-paste from autograd to torch * little fix in torch.py; get tests passing * changelog * put back the default tensor type after tests --------- Co-authored-by: Christina Lee <christina@xanadu.ai>

* First version of StateVectorProjector * Outer product fix * Support projector in default qubit * change inheritance * consistent diagonalizing gates * Update docstring * Proper bra-ket in label * Tests for StateVectorProjector * fix tests * Projector prototype * Explicit signature and pow method * discard hacky prototype * New hack (thanks Tymmy <3) * Exception tests * Docstring for __new__ * Label method adds matrix to cache * tests for label method * Update docstrings * Update changelog * enhanced docstring for label method * `Projector.__new__` docstring rephrased * Fix docstring Co-authored-by: Matthew Silverman <matthews@xanadu.ai> * Remove boolean kwarg * Remove `basis_representation` from changelog * Fix wire length issue * Update Projector example in docstring Co-authored-by: Matthew Silverman <matthews@xanadu.ai> * Update changelog description Co-authored-by: Tom Bromley <49409390+trbromley@users.noreply.github.com> * Update projector description Co-authored-by: Tom Bromley <49409390+trbromley@users.noreply.github.com> * change argument names to be `state` * Remove shape indication in docstring Co-authored-by: Tom Bromley <49409390+trbromley@users.noreply.github.com> * fix expval with state vector projector * Code example and remove hidden class docstring * minor docstring corrections * add expval test (it was failing) * projector bind new parameters dispatcher * Update error string Co-authored-by: Matthew Silverman <matthews@xanadu.ai> * fix projector and qubit device * extensive projector testing * remove outdated test * fix tests typo * Additional indications on input shape Co-authored-by: Mudit Pandey <mudit.pandey@xanadu.ai> * revision feedback * update copy --------- Co-authored-by: = <=> Co-authored-by: Matthew Silverman <matthews@xanadu.ai> Co-authored-by: Tom Bromley <49409390+trbromley@users.noreply.github.com> Co-authored-by: Mudit Pandey <mudit.pandey@xanadu.ai>

* update test * fix comments * see if this fixes mocking * Add pytest fixture Hermitian._eigs={} before each test * Update tests/ops/qubit/test_observables.py Co-authored-by: Matthew Silverman <matthews@xanadu.ai> --------- Co-authored-by: Matthew Silverman <matthews@xanadu.ai>

Co-authored-by: soranjh <40344468+soranjh@users.noreply.github.com>

* save changes to tapering * add tapering update + tests * more testing * lint and tests * more tests * Apply suggestions from code review Co-authored-by: soranjh <40344468+soranjh@users.noreply.github.com> * added more tests * fix bug * lint * lint * lint * lint + codecov * more tests * Added test and warning message for coverage * Update pennylane/ops/functions/generator.py Co-authored-by: Utkarsh <utkarshazad98@gmail.com> * lint * lint and codefactor * lint + changelog * fix tests + lint * code review --------- Co-authored-by: soranjh <40344468+soranjh@users.noreply.github.com> Co-authored-by: Utkarsh <utkarshazad98@gmail.com>

* switch len with qml.math.size * update change log

* Add deprecation warning to public methods * Updated to not raise warnings in normal execution * Updated deprecation page * Added tests * Added qnode execution test * Updated tests to use recwarn * Added moer context to deprecations * Updated deprecations; removed warnings * Reverted pytest.ini * Removed unused import * Apply suggestions from code review Co-authored-by: Tom Bromley <49409390+trbromley@users.noreply.github.com> * Updated docstring * Added more warnings * Fixed name * Reformatting * Reverted changes to `operations`, `observables` * Updated where deprecation is shown * Formatting * Updated plugin doc * Updated per review --------- Co-authored-by: Tom Bromley <49409390+trbromley@users.noreply.github.com>

) * Updated `bind_new_parameters` * Updated `bind_new_parameters` and tests * New op * Added tests * Updated changelog * Black * Fixed error * Updated for CodeFactor

* Support broadcasting in state measurements * docs for is_state_batched * Support broadcasting in sample measurements * Support broadcastingin simulate and preprocess --------- Co-authored-by: Matthew Silverman <matthews@xanadu.ai>

* Support classical shadow measurements * Add shot vector tests * Address comments * Add comment for diag_list * Apply suggestions from code review Co-authored-by: Christina Lee <christina@xanadu.ai> * Fix * Add unit tests for process_state_with_shots * Update preprocessing * Apply suggestions from code review Co-authored-by: Frederik Wilde <42576579+frederikwilde@users.noreply.github.com> * Address PR comments * set rng for test --------- Co-authored-by: Christina Lee <christina@xanadu.ai> Co-authored-by: Frederik Wilde <42576579+frederikwilde@users.noreply.github.com>

* support trainable Sum observables (analytic only) * just use pre-rotated state; use super with finite shots * fix tests; add test for trainable Sum coeffs * changelog * use overlapping wires to prove they work * add hacky gradient support * Revert "add hacky gradient support" This reverts commit 09009c8. * set interface=None to keep test behaviour

* integrate qnode with new device * some diff method improvements * repr methods * add tests, always pass config to device * add tests, always pass config to device * final test * pylint * autograd integration tests * pylint * Update pennylane/interfaces/execution.py Co-authored-by: Matthew Silverman <matthews@xanadu.ai> * pass shots through methods * changelog * revert set shots change * Apply suggestions from code review Co-authored-by: Matthew Silverman <matthews@xanadu.ai> * revert executionc hange, pylint: * pylint again --------- Co-authored-by: Matthew Silverman <matthews@xanadu.ai> Co-authored-by: Romain Moyard <rmoyard@gmail.com>

* pre-release for v0.31.0 * Update doc/releases/changelog-0.31.0.md

* pre-release for v0.31.0 * Update doc/releases/changelog-0.31.0.md * Incrementing the version number to v0.32.0-dev

Update Python version to 3.8 Co-authored-by: Romain Moyard <rmoyard@gmail.com>

* Support `HardwareHamiltonian` pulses in `stoch_pulse_grad` (#4215) * single out gradient transform checks * rename stochastic pulse gradient file * unify gradient_analysis and grad_method_validation * continue restructure of analysis+validation * CV * black * modularize more * more modularizing * black * tiny [skip ci] * [skip ci] lint * remove dummy test * test fix * add test file to linting test file * test fixes, docstrings * code review * docstring gradient_analysis_and_grad_method_validation * move first fun * code review:move functions * test regex * regexs * move and promote reorder_grads * tmp * more tmp * test cases, contractions * lint * docstring * even more tmp * cleanup * black * tmp * lint * move stoch_pulse_gradient.. files back to pulse_gradient... * move stoch_pulse_gradient.. files back to pulse_gradient... * lint * rename * extend functions and tests * lint and black * changelog * improve * update example to include non-Pauli word generator * add jit test with pauli sentence * tmp * debugging, docstring, extend test * review * optimize for Pauli words * Apply suggestions from code review Co-authored-by: Romain Moyard <rmoyard@gmail.com> * test cases code review * fix parametrization * drafting * working prototype * finish merge; cleanup * changelog * comments * [skip ci] * raising an error; cleanup [skip ci] * Apply suggestions from code review Co-authored-by: Korbinian Kottmann <43949391+Qottmann@users.noreply.github.com> * change contraction idea * typo in docs * tests * remove prints * fix test * test descriptions * fix merge * format * code review; test coverage * coverage reordering * fix * trigger CI * trigger * clear caches * trigger * trigger --------- Co-authored-by: Romain Moyard <rmoyard@gmail.com> Co-authored-by: Korbinian Kottmann <43949391+Qottmann@users.noreply.github.com> Co-authored-by: Korbinian Kottmann <Korbinian.Kottmann@gmail.com> * Fix batching of derivative tapes in autograd (#4245) * exclude files from pr --------- Co-authored-by: David Wierichs <david.wierichs@xanadu.ai> Co-authored-by: Romain Moyard <rmoyard@gmail.com> Co-authored-by: Korbinian Kottmann <43949391+Qottmann@users.noreply.github.com> Co-authored-by: Korbinian Kottmann <Korbinian.Kottmann@gmail.com> Co-authored-by: Christina Lee <christina@xanadu.ai> Co-authored-by: GitHub Actions Bot <>

…nd ndim attribute.

…hape and for performance increase in multi_dispatch.unwrap

…autoray.

…attribute.

trbromley · 2023-07-17T13:06:49Z

[SC-40363]

codecov · 2023-07-17T16:50:20Z

Codecov Report

Merging #4261 (eeccf62) into master (bb74af8) will increase coverage by 0.00%.
The diff coverage is 100.00%.

@@           Coverage Diff           @@
##           master    #4261   +/-   ##
=======================================
  Coverage   99.77%   99.77%           
=======================================
  Files         351      351           
  Lines       32498    32510   +12     
=======================================
+ Hits        32425    32437   +12     
  Misses         73       73

Impacted Files	Coverage Δ
pennylane/math/__init__.py	`100.00% <100.00%> (ø)`
pennylane/math/multi_dispatch.py	`100.00% <100.00%> (ø)`
pennylane/math/single_dispatch.py	`99.72% <100.00%> (-0.01%)`	⬇️

vincentmr

Looks good to me, thanks for the nice work @frederikwilde !

Remove custom registration of ndim and shape and use caching for in-p…

2b608b6

…lace import module import.

frederikwilde added the WIP 🚧 Work-in-progress label Jun 16, 2023

frederikwilde and others added 26 commits June 16, 2023 17:55

Implement custom ndim and shape in qml.math to catch slow autograd Ar…

1cfba84

…rayBox operations. Disable _check_batching for set_parameters in copied operators to reduce do overheads.

Add new string representation to FermiWord (#4255)

054c93e

Update Jax version in the CI 0.4.12 (#4259)

58c45db

torch integration for DefaultQubit2 (#4257)

5dd87fe

* copy-paste from autograd to torch * little fix in torch.py; get tests passing * changelog * put back the default tensor type after tests --------- Co-authored-by: Christina Lee <christina@xanadu.ai>

move dispatching into function to fix docstrings (#4260)

86ab67e

Co-authored-by: soranjh <40344468+soranjh@users.noreply.github.com>

Fix bug in fermi op arithmetic for np.array([1, 2])[0] (#4262)

ddc8ef8

* switch len with qml.math.size * update change log

[skip ci] Add more registered dispatches to bind_new_parameters (#4256

af7c355

) * Updated `bind_new_parameters` * Updated `bind_new_parameters` and tests * New op * Added tests * Updated changelog * Black * Fixed error * Updated for CodeFactor

Add the fermi module docs (#4250)

7f96eae

pre-release for v0.31.0 (#4265)

392c7d5

* pre-release for v0.31.0 * Update doc/releases/changelog-0.31.0.md

Incrementing the version number to v0.32.0-dev (#4266)

93ee5bb

* pre-release for v0.31.0 * Update doc/releases/changelog-0.31.0.md * Incrementing the version number to v0.32.0-dev

Update installation.rst (#4247)

b103aeb

Update Python version to 3.8 Co-authored-by: Romain Moyard <rmoyard@gmail.com>

add emojis to changelog (#4267)

0ff6c3e

frederikwilde and others added 3 commits July 7, 2023 10:55

Remove custom ndim and shape for ArrayBoxes since they have a shape a…

e99ba1a

…nd ndim attribute.

Add custom ArrayBox handling for demonstration purposes in ndim and s…

42b5b90

…hape and for performance increase in multi_dispatch.unwrap

Merge branch 'master' into slowdowns-autoray

a90969c

frederikwilde added do not merge ⚠️ Do not merge the pull request until this label is removed and removed WIP 🚧 Work-in-progress labels Jul 10, 2023

frederikwilde added 2 commits July 13, 2023 12:16

resolve merge conflict

10305ef

Resolve merge error.

b219291

vincentmr mentioned this pull request Jul 14, 2023

Add multiprocessing to DQ2 #4319

Merged

5 tasks

frederikwilde added 4 commits July 14, 2023 10:29

Clean up

c0cdb2b

Merge branch 'master' into slowdowns-autoray

ed6db02

Make ndim and shape first check for attributes before dispatching to …

94401f4

…autoray.

Remove test for jax.numpy.ndim, since we will always access the ndim …

4050321

…attribute.

frederikwilde added 5 commits July 17, 2023 11:00

Fix torch error

3dab5d5

Merge branch 'master' into slowdowns-autoray

558f3f1

Black

cc5ca36

Changelog [skip-ci]

3a1fca2

Changelog

4f04560

frederikwilde removed the do not merge ⚠️ Do not merge the pull request until this label is removed label Jul 17, 2023

Remove unused import

67fe086

frederikwilde and others added 3 commits July 17, 2023 15:38

Remove unreachable code.

acc5c76

Merge branch 'master' into slowdowns-autoray

b8cea89

Merge branch 'master' into slowdowns-autoray

eeccf62

vincentmr self-requested a review July 25, 2023 17:14

vincentmr approved these changes Jul 25, 2023

View reviewed changes

mudit2812 added this to the v0.32 milestone Aug 14, 2023

trbromley removed this from the v0.32 milestone Aug 17, 2023

frederikwilde closed this Aug 23, 2023

trbromley deleted the slowdowns-autoray branch March 7, 2024 13:49

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Speed up dispatching with autoray #4261

Speed up dispatching with autoray #4261

frederikwilde commented Jun 16, 2023 •

edited

Loading

frederikwilde commented Jun 16, 2023

github-actions bot commented Jun 16, 2023

trbromley commented Jul 17, 2023

codecov bot commented Jul 17, 2023 •

edited

Loading

vincentmr left a comment

Speed up dispatching with autoray #4261

Speed up dispatching with autoray #4261

Conversation

frederikwilde commented Jun 16, 2023 • edited Loading

frederikwilde commented Jun 16, 2023

github-actions bot commented Jun 16, 2023

trbromley commented Jul 17, 2023

codecov bot commented Jul 17, 2023 • edited Loading

Codecov Report

vincentmr left a comment

Choose a reason for hiding this comment

frederikwilde commented Jun 16, 2023 •

edited

Loading

codecov bot commented Jul 17, 2023 •

edited

Loading