Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix the issue with JAX JIT for grad when diff_method="adjoint" and mode="backward" #2460

Merged
merged 6 commits into from Apr 20, 2022

Conversation

maliasadi
Copy link
Member

Context:
Fix issue #2449 (and add make format).

@github-actions
Copy link
Contributor

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.

@codecov
Copy link

codecov bot commented Apr 19, 2022

Codecov Report

❗ No coverage uploaded for pull request base (v0.23.0-rc0@82da18c). Click here to learn what that means.
The diff coverage is n/a.

@@              Coverage Diff               @@
##             v0.23.0-rc0    #2460   +/-   ##
==============================================
  Coverage               ?   99.47%           
==============================================
  Files                  ?      244           
  Lines                  ?    19453           
  Branches               ?        0           
==============================================
  Hits                   ?    19351           
  Misses                 ?      102           
  Partials               ?        0           

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 82da18c...c02ecd2. Read the comment docs.

Copy link
Contributor

@antalszava antalszava left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you for fixing this @maliasadi 🎉 Looks good, just curious about the use of doctest in pennylane/interfaces/jax_jit.py.

pennylane/interfaces/jax_jit.py Outdated Show resolved Hide resolved
Makefile Show resolved Hide resolved
pennylane/interfaces/jax_jit.py Outdated Show resolved Hide resolved
pennylane/interfaces/jax_jit.py Outdated Show resolved Hide resolved
Copy link
Contributor

@antalszava antalszava left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks good 🎉

@antalszava
Copy link
Contributor

@maliasadi the conflict comes from the fact that doc/releases/changelog-dev.md has recently been renamed to doc/releases/changelog-0.23.0.md.

Copy link
Contributor

@rmoyard rmoyard left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It looks good to me 💯

@@ -472,6 +471,31 @@ def cost(a, b):
expected = -2 * np.sin(2 * a)
assert np.allclose(jac, expected, atol=tol, rtol=0)

def test_jit_grad_with_backward_mode(self, execute_kwargs, interface):
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Great test!


shapes = [jax.ShapeDtypeStruct((1, len(p)), dtype) for p in params]
jacs = host_callback.call(jacs_wrapper, params, result_shape=shapes)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Clever solution! 💯

@antalszava
Copy link
Contributor

Locally running pylint checks out. 👍

@antalszava antalszava merged commit 1bab718 into v0.23.0-rc0 Apr 20, 2022
@antalszava antalszava deleted the fix_issue_2449 branch April 20, 2022 15:43
antalszava added a commit that referenced this pull request Apr 22, 2022
* Remove deprecated `get_unitary_matrix` (#2457)

* Remove

* Changelog

* changelog: breaking change

* Trigger CI

Co-authored-by: Antal Szava <antalszava@gmail.com>

* Rename `doc/releases/changelog-dev.md` → `doc/releases/changelog-0.23.0.md` (release candidate) (#2461)

* changelog rename

* update ref of the file

* Fix the issue with JAX JIT for grad when `diff_method="adjoint"` and `mode="backward"` (#2460)

* Fix the issue

* Apply code review suggestions

* Update changelog

* Update jax.py

* Trigger CI

* Bugfix for Nesterov Momentum Optimier (#2466)

* nesterov_fix

* nesterov momentum bugfix

* Update doc/releases/changelog-0.23.0.md

* Update tests/optimize/test_optimize.py

* Use pennylane numpy in molecular hamiltonian (#2465)

* fix pl numpy and ham complex coeff

* test

* PR num

* test tidy

* correct method for returning real coeffs

* linting

Co-authored-by: Antal Szava <antalszava@gmail.com>

* Remove `qml.finite_diff` (#2464)

* remove finite-diff

* remove unused import

* Update doc/releases/changelog-0.23.0.md

Co-authored-by: antalszava <antalszava@gmail.com>

* remove hamiltonian diff test file

Co-authored-by: antalszava <antalszava@gmail.com>

* Update qchem documentation (#2454)

* exclude files from pr

* exclude more files

Co-authored-by: Romain Moyard <rmoyard@gmail.com>
Co-authored-by: Antal Szava <antalszava@gmail.com>
Co-authored-by: Ali Asadi <ali@xanadu.ai>
Co-authored-by: Christina Lee <christina@xanadu.ai>
Co-authored-by: soranjh <40344468+soranjh@users.noreply.github.com>
Co-authored-by: GitHub Actions Bot <>
Co-authored-by: dime10 <dime10@users.noreply.github.com>
dime10 added a commit that referenced this pull request Apr 22, 2022
* Remove deprecated `get_unitary_matrix` (#2457)

* Remove

* Changelog

* changelog: breaking change

* Trigger CI

Co-authored-by: Antal Szava <antalszava@gmail.com>

* Rename `doc/releases/changelog-dev.md` → `doc/releases/changelog-0.23.0.md` (release candidate) (#2461)

* changelog rename

* update ref of the file

* Fix the issue with JAX JIT for grad when `diff_method="adjoint"` and `mode="backward"` (#2460)

* Fix the issue

* Apply code review suggestions

* Update changelog

* Update jax.py

* Trigger CI

* Bugfix for Nesterov Momentum Optimier (#2466)

* nesterov_fix

* nesterov momentum bugfix

* Update doc/releases/changelog-0.23.0.md

* Update tests/optimize/test_optimize.py

* Use pennylane numpy in molecular hamiltonian (#2465)

* fix pl numpy and ham complex coeff

* test

* PR num

* test tidy

* correct method for returning real coeffs

* linting

Co-authored-by: Antal Szava <antalszava@gmail.com>

* Remove `qml.finite_diff` (#2464)

* remove finite-diff

* remove unused import

* Update doc/releases/changelog-0.23.0.md

Co-authored-by: antalszava <antalszava@gmail.com>

* remove hamiltonian diff test file

Co-authored-by: antalszava <antalszava@gmail.com>

* Update qchem documentation (#2454)

* fix mera formatting issues (#2470)

* Fix qcut doc render (#2471)

* Update

* Update

* Trigger CI

* fix has_matrix (#2476)

* use qml.eigvals because MeasurementProcess.eigvals is being deprecated and raises a warning (#2477)

* exclude files from pr

* exclude more files

Co-authored-by: Romain Moyard <rmoyard@gmail.com>
Co-authored-by: Antal Szava <antalszava@gmail.com>
Co-authored-by: Ali Asadi <ali@xanadu.ai>
Co-authored-by: Christina Lee <christina@xanadu.ai>
Co-authored-by: soranjh <40344468+soranjh@users.noreply.github.com>
Co-authored-by: GitHub Actions Bot <>
Co-authored-by: dime10 <dime10@users.noreply.github.com>
antalszava added a commit that referenced this pull request Apr 25, 2022
* Remove deprecated `get_unitary_matrix` (#2457)

* Remove

* Changelog

* changelog: breaking change

* Trigger CI

Co-authored-by: Antal Szava <antalszava@gmail.com>

* Rename `doc/releases/changelog-dev.md` → `doc/releases/changelog-0.23.0.md` (release candidate) (#2461)

* changelog rename

* update ref of the file

* Fix the issue with JAX JIT for grad when `diff_method="adjoint"` and `mode="backward"` (#2460)

* Fix the issue

* Apply code review suggestions

* Update changelog

* Update jax.py

* Trigger CI

* Bugfix for Nesterov Momentum Optimier (#2466)

* nesterov_fix

* nesterov momentum bugfix

* Update doc/releases/changelog-0.23.0.md

* Update tests/optimize/test_optimize.py

* Use pennylane numpy in molecular hamiltonian (#2465)

* fix pl numpy and ham complex coeff

* test

* PR num

* test tidy

* correct method for returning real coeffs

* linting

Co-authored-by: Antal Szava <antalszava@gmail.com>

* Remove `qml.finite_diff` (#2464)

* remove finite-diff

* remove unused import

* Update doc/releases/changelog-0.23.0.md

Co-authored-by: antalszava <antalszava@gmail.com>

* remove hamiltonian diff test file

Co-authored-by: antalszava <antalszava@gmail.com>

* Update qchem documentation (#2454)

* fix mera formatting issues (#2470)

* Fix qcut doc render (#2471)

* Update

* Update

* Trigger CI

* fix has_matrix (#2476)

* use qml.eigvals because MeasurementProcess.eigvals is being deprecated and raises a warning (#2477)

* Warn upon import of an `ObservableReturnType` from operation instead of measurements (#2481)

* warn upon import of observablereturntype from operation

* black?

* Update pennylane/__init__.py

Co-authored-by: antalszava <antalszava@gmail.com>

* Update tests/test_operation.py

Co-authored-by: antalszava <antalszava@gmail.com>

* update changelog

* Update pennylane/operation.py

Co-authored-by: antalszava <antalszava@gmail.com>

* Update chemistry doc (#2482)

* Automatic upload to pip when released (#2484)

* upload

* Return line

* Jacobian Slice Bug StrawberryFields (#2485)

* fixed jacobian slicing issue

Update pennylane/interfaces/autograd.py

Co-authored-by: antalszava <antalszava@gmail.com>

* added testcase for custom device

* Update doc/releases/changelog-0.23.0.md

* retrigger checks

Co-authored-by: antalszava <antalszava@gmail.com>

* Documentation fixes for the v0.23.0 release (#2472)

* Fix qml.ctrl docstring

* Add notice about control values other than 0,1

* Fix bell example on measurements docs page

* Update doc/introduction/measurements.rst

Co-authored-by: Christina Lee <christina@xanadu.ai>

* Update pennylane/transforms/control.py

Co-authored-by: Christina Lee <christina@xanadu.ai>

* Make the cut_circuit_mc docs example deterministic

* Fix sphinx build

Co-authored-by: Christina Lee <christina@xanadu.ai>
Co-authored-by: antalszava <antalszava@gmail.com>

* Issue a warning about repeated cached executions with shot noise (#2478)

* Add warning about caching with finite shots

* Fix import order

* Add changelog

* Use set_shots to determine shots for warning

* format

* disable pylint 'protected-access'

* Fixes and cleanup to the inspection logic

* format

* add test for warning

* Review: expand warning message

* Review: add additional tests

Co-authored-by: antalszava <antalszava@gmail.com>

* exclude files from pr

* Exclude more files

Co-authored-by: Romain Moyard <rmoyard@gmail.com>
Co-authored-by: Antal Szava <antalszava@gmail.com>
Co-authored-by: Ali Asadi <ali@xanadu.ai>
Co-authored-by: Christina Lee <christina@xanadu.ai>
Co-authored-by: soranjh <40344468+soranjh@users.noreply.github.com>
Co-authored-by: Sam Banning <samcbanning@gmail.com>
Co-authored-by: David Ittah <dime10@users.noreply.github.com>
Co-authored-by: GitHub Actions Bot <>
antalszava added a commit that referenced this pull request Apr 26, 2022
* Remove deprecated `get_unitary_matrix` (#2457)

* Remove

* Changelog

* changelog: breaking change

* Trigger CI

Co-authored-by: Antal Szava <antalszava@gmail.com>

* Rename `doc/releases/changelog-dev.md` → `doc/releases/changelog-0.23.0.md` (release candidate) (#2461)

* changelog rename

* update ref of the file

* Fix the issue with JAX JIT for grad when `diff_method="adjoint"` and `mode="backward"` (#2460)

* Fix the issue

* Apply code review suggestions

* Update changelog

* Update jax.py

* Trigger CI

* Bugfix for Nesterov Momentum Optimier (#2466)

* nesterov_fix

* nesterov momentum bugfix

* Update doc/releases/changelog-0.23.0.md

* Update tests/optimize/test_optimize.py

* Use pennylane numpy in molecular hamiltonian (#2465)

* fix pl numpy and ham complex coeff

* test

* PR num

* test tidy

* correct method for returning real coeffs

* linting

Co-authored-by: Antal Szava <antalszava@gmail.com>

* Remove `qml.finite_diff` (#2464)

* remove finite-diff

* remove unused import

* Update doc/releases/changelog-0.23.0.md

Co-authored-by: antalszava <antalszava@gmail.com>

* remove hamiltonian diff test file

Co-authored-by: antalszava <antalszava@gmail.com>

* Update qchem documentation (#2454)

* fix mera formatting issues (#2470)

* Fix qcut doc render (#2471)

* Update

* Update

* Trigger CI

* fix has_matrix (#2476)

* use qml.eigvals because MeasurementProcess.eigvals is being deprecated and raises a warning (#2477)

* Warn upon import of an `ObservableReturnType` from operation instead of measurements (#2481)

* warn upon import of observablereturntype from operation

* black?

* Update pennylane/__init__.py

Co-authored-by: antalszava <antalszava@gmail.com>

* Update tests/test_operation.py

Co-authored-by: antalszava <antalszava@gmail.com>

* update changelog

* Update pennylane/operation.py

Co-authored-by: antalszava <antalszava@gmail.com>

* Update chemistry doc (#2482)

* Automatic upload to pip when released (#2484)

* upload

* Return line

* Jacobian Slice Bug StrawberryFields (#2485)

* fixed jacobian slicing issue

Update pennylane/interfaces/autograd.py

Co-authored-by: antalszava <antalszava@gmail.com>

* added testcase for custom device

* Update doc/releases/changelog-0.23.0.md

* retrigger checks

Co-authored-by: antalszava <antalszava@gmail.com>

* Documentation fixes for the v0.23.0 release (#2472)

* Fix qml.ctrl docstring

* Add notice about control values other than 0,1

* Fix bell example on measurements docs page

* Update doc/introduction/measurements.rst

Co-authored-by: Christina Lee <christina@xanadu.ai>

* Update pennylane/transforms/control.py

Co-authored-by: Christina Lee <christina@xanadu.ai>

* Make the cut_circuit_mc docs example deterministic

* Fix sphinx build

Co-authored-by: Christina Lee <christina@xanadu.ai>
Co-authored-by: antalszava <antalszava@gmail.com>

* Issue a warning about repeated cached executions with shot noise (#2478)

* Add warning about caching with finite shots

* Fix import order

* Add changelog

* Use set_shots to determine shots for warning

* format

* disable pylint 'protected-access'

* Fixes and cleanup to the inspection logic

* format

* add test for warning

* Review: expand warning message

* Review: add additional tests

Co-authored-by: antalszava <antalszava@gmail.com>

* Fix QChem documentation (#2487)

* Fix finite-shots/caching warning appearing internally for gradients (#2491)

* Fix warning appearing internally for gradients

* pylint: fix protected attribute

* Add test for checking no warning is raised

* V0.23.0 notes (#2474)

* notes + version

* sections

* update examples

* remove empty sections

* has_matrix bug fix

* Adding more PRs to the changelog

* PR, extend contributors

* Extend list of contributors

* Update the order of entries

* Add auto cutting example; reorg QChem

* remove curr release from 0.22.2 header

* formatting

* fix link

* Update

* Trigger CI

* Change code blooc

* Typo

* Update pennylane/transforms/optimization/pattern_matching.py

Co-authored-by: antalszava <antalszava@gmail.com>

* Update pennylane/transforms/optimization/pattern_matching.py

Co-authored-by: antalszava <antalszava@gmail.com>

* Update pennylane/transforms/optimization/pattern_matching.py

Co-authored-by: antalszava <antalszava@gmail.com>

* Update pennylane/transforms/optimization/pattern_matching.py

Co-authored-by: antalszava <antalszava@gmail.com>

* Addd seealso

* modify the qchem section

* Update doc/releases/changelog-0.23.0.md

Co-authored-by: Josh Izaac <josh146@gmail.com>

* Update doc/releases/changelog-0.23.0.md

Co-authored-by: Josh Izaac <josh146@gmail.com>

* Update doc/releases/changelog-0.23.0.md

Co-authored-by: Josh Izaac <josh146@gmail.com>

* Update doc/releases/changelog-0.23.0.md

Co-authored-by: Josh Izaac <josh146@gmail.com>

* code review

* code review

* code review (remove empty lines & add pattern matching opt docstring link)

* Update doc/releases/changelog-0.23.0.md

Co-authored-by: Josh Izaac <josh146@gmail.com>

* Update doc/releases/changelog-0.23.0.md

Co-authored-by: Josh Izaac <josh146@gmail.com>

* shape and numeric type example

* MP description update

* rephrase expand_strategy for qml.specs improvement

* update

* caching phrasing

* removed

* phrasing

* updates

* bug phrasing, reorder

* update link placement and use pycon

* include swap test output in example

* More info on HS template

* rephrase QCut

* move example to UsageDetails

* auto QCut rephrase

* More powerful circuit cutting

* Trigger CI

* Upddate

* add qchem code example

* correct typo in qchem text

* no backticks

* fix the case when obs=None for a MeasurementProcess for interface selection

* update

* changelog

* docstring

* Update doc/releases/changelog-0.23.0.md

Co-authored-by: Josh Izaac <josh146@gmail.com>

* update links as per code review

* typo

* pycon syntax

* double backticks

* no dev in release header

* Lightning v0.23.0 required in setup.py

Co-authored-by: Romain Moyard <rmoyard@gmail.com>
Co-authored-by: soranjh <soranjh@yahoo.com>
Co-authored-by: Josh Izaac <josh146@gmail.com>

* Remove typo

* Update pennylane/interfaces/jax.py

Co-authored-by: Romain Moyard <rmoyard@gmail.com>
Co-authored-by: Ali Asadi <ali@xanadu.ai>
Co-authored-by: Christina Lee <christina@xanadu.ai>
Co-authored-by: soranjh <40344468+soranjh@users.noreply.github.com>
Co-authored-by: Sam Banning <samcbanning@gmail.com>
Co-authored-by: David Ittah <dime10@users.noreply.github.com>
Co-authored-by: Josh Izaac <josh146@gmail.com>
Co-authored-by: soranjh <soranjh@yahoo.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

None yet

3 participants