Fixing JAX interface for adjoint diff #1349

albi3ro · 2021-05-21T14:28:14Z

Closes Issue #1345

Currently, the jax interface fails to return correct results with the adjoint differentiation method. The Jax interface makes a copy of the current tape to take it's derivative, but the tape copy didn't copy over differentiation information.

I have attempted to rectify this by improving the tape copy method, but now trying to take the derivative of the circuit causes my kernel to end.

github-actions · 2021-05-21T14:28:33Z

Hello. You may have forgotten to update the changelog!
Please edit .github/CHANGELOG.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.

pennylane/tape/jacobian_tape.py

codecov · 2021-05-21T15:05:56Z

Codecov Report

Merging #1349 (97f23b8) into master (23b194e) will increase coverage by 0.00%.
The diff coverage is 100.00%.

@@           Coverage Diff           @@
##           master    #1349   +/-   ##
=======================================
  Coverage   98.16%   98.16%           
=======================================
  Files         154      154           
  Lines       11544    11548    +4     
=======================================
+ Hits        11332    11336    +4     
  Misses        212      212

Impacted Files	Coverage Δ
pennylane/interfaces/autograd.py	`100.00% <ø> (ø)`
pennylane/tape/jacobian_tape.py	`98.00% <100.00%> (+0.04%)`	⬆️

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 23b194e...97f23b8. Read the comment docs.

albi3ro · 2021-05-21T18:59:49Z

pennylane/qnode.py

+            jnp.array(0.1, dtype=jnp.float64)
+        except UserWarning:
+            return False
+    return True


I don't know how to test this line, given float64 mode must be configured on startup.

josh146

Thanks for catching this @albi3ro! Just two suggestions re: improving the precision check and moving the location of the warning logic, but once addressed this should be good for merging.

.github/CHANGELOG.md

josh146 · 2021-05-26T05:21:44Z

pennylane/qnode.py

+    import jax.numpy as jnp
+
+    with warnings.catch_warnings():
+        warnings.filterwarnings("error")
+        try:
+            jnp.array(0.1, dtype=jnp.float64)
+        except UserWarning:
+            return False
+    return True


Nice workaround, but it feels a bit too reliant on implementation details of JAX and how it raises warnings.

I would suggest instead simply querying the JAX configuration:

>>> from jax.config import config >>> config.x64_enabled False >>> config.update("jax_enable_x64", True) >>> config.x64_enabled True

Also, I would recommend moving this to interfaces/jax.py, so that all JAX imports are isolated in a single place. The warning can then be raised by the tape mixin class querying jacobian_options dictionary

josh146 · 2021-05-26T05:29:03Z

pennylane/qnode.py

+        if interface == "jax":
+            if not _jax_float64_support():
+                warnings.warn(
+                    "float64 support not enabled for jax. "
+                    "May cause inaccuracies for `diff_method='finite-diff'`",
+                    UserWarning,
+                )
+
+        return JacobianTape, interface, device, {"method": "numeric"}


Yep, I would say this definitely feels like something that lives on the interface rather than the QNode :) Especially because it simply raises a warning, rather than changing the diff method.

albi3ro · 2021-05-27T20:02:11Z

.github/CHANGELOG.md

 * Fixes drawing QNodes that contain multiple measurements on a single wire.
  [(#1353)](https://github.com/PennyLaneAI/pennylane/pull/1353)

+


Suggested change

antalszava · 2021-05-28T19:13:03Z

pennylane/tape/jacobian_tape.py

@@ -533,7 +538,8 @@ def jacobian(self, device, params=None, **options):
                # First order (forward) finite-difference will be performed.
                # Compute the value of the tape at the current parameters here. This ensures
                # this computation is only performed once, for all parameters.
-                options["y0"] = np.asarray(self.execute_device(params, device))
+                params_f64 = np.array(params, dtype=np.float64)


One of the jax device tests doesn't seem to like this line 🤔
https://github.com/PennyLaneAI/pennylane/pull/1349/checks?check_run_id=2697038221#step:10:63

On another note: wouldn't this casting be a bit too specific at a general place? This would potentially cast the parameters all the time regardless of the interface. Would we only like to do it for the jax interface?

@albi3ro I think I agree with @antalszava in this case. While there is a good argument for enforcing finite difference to always use float64, I feel like this should be an interface level option/default, rather than enforced in the Jacobian method

antalszava · 2021-05-28T19:14:06Z

tests/tape/test_jacobian_tape.py

@@ -55,6 +55,23 @@ def test_parameter_info(self, make_tape):
        }


+class TestTapeCopying:


Might be worth adding an integration test similar to the original example in the #1345 issue to ensure that the results using the jax interface are correct.

antalszava · 2021-05-28T19:16:10Z

Hi @albi3ro, it must have been quite some digging to work this out, nice one! 🥇 One device test seems to be failing and also wondering if as Josh suggested, it could be good to add the fix specifically to the jax interface code.

albi3ro · 2021-05-28T19:27:49Z

The one test is failing because of an issue I just raised, Issue #1369 . When a parameter is a matrix instead of float, params is a deprecated ragged array and can't be cast to float64.

While this problem arose with the jax interface, it will actually occur any time we use finite differences with float32, or any other less precise data type. Therefore, I think the solution should be more general than just the jax interface.

numeric_pd uses float64 parameters anyway. "y0" just uses the less precise data type. Simple fix, just one that took me forever to find.

josh146

Thanks @albi3ro!

When a parameter is a matrix instead of float, params is a deprecated ragged array and can't be cast to float64.

I think this might be a bug with the test not specifying requires_grad=False on the matrix, as opposed to a bug in the code, see
#1369 (comment)

I have two requests:

I feel like the general float64 fix within the JacobianTape is out of scope for this PR, and should be considered separately. Instead, I think it should be sufficient to just make the JAX interface use float64 in the finite-diff tests (with maybe small modifications to jax_ops.py).
Does this mean that there are no adjoint tests for the JAX interface? We should probably add them now to ensure that this bug is fixed :)

josh146 · 2021-05-31T07:43:29Z

pennylane/tape/jacobian_tape.py

@@ -533,7 +538,8 @@ def jacobian(self, device, params=None, **options):
                # First order (forward) finite-difference will be performed.
                # Compute the value of the tape at the current parameters here. This ensures
                # this computation is only performed once, for all parameters.
-                options["y0"] = np.asarray(self.execute_device(params, device))
+                params_f64 = np.array(params, dtype=np.float64)


@albi3ro I think I agree with @antalszava in this case. While there is a good argument for enforcing finite difference to always use float64, I feel like this should be an interface level option/default, rather than enforced in the Jacobian method

albi3ro · 2021-05-31T13:43:59Z

I added the finite-diff change to this PR because I thought it was going to be easy. But there needs to better a better solution than just the easy one, so that is getting moved to it's own PR.

I had the adjoint test before, but it got removed when I fire and brimstoned the PR after a bad merge that messed everything up. It's back. Everything should be good to go now.

josh146

Looks good @albi3ro 🙂

josh146 · 2021-05-31T15:43:30Z

pennylane/interfaces/autograd.py

@@ -119,7 +119,7 @@ def get_parameters(self, trainable_only=True, return_arraybox=False):
                qml.RY(0.543, wires=0)
                qml.CNOT(wires=[0, 'a'])
                qml.RX(0.133, wires='a')
-                expval(qml.PauliZ(wires=[0]))
+                qml.expval(qml.PauliZ(wires=[0]))


oops, good catch

josh146 · 2021-05-31T15:47:39Z

tests/test_qubit_device_adjoint_jacobian.py

+        grad_adjoint = jax.grad(qnode_adjoint)(params1, params2)
+        grad_backprop = jax.grad(qnode_backprop)(params1, params2)
+
+        assert np.allclose(grad_adjoint, grad_backprop)


Good to have this test in 💪

josh146 reviewed May 21, 2021

View reviewed changes

pennylane/tape/jacobian_tape.py Outdated Show resolved Hide resolved

josh146 reviewed May 21, 2021

View reviewed changes

pennylane/tape/jacobian_tape.py Outdated Show resolved Hide resolved

albi3ro requested a review from josh146 May 21, 2021 18:57

albi3ro commented May 21, 2021

View reviewed changes

albi3ro added the review-ready 👌 PRs which are ready for review by someone from the core team. label May 21, 2021

albi3ro requested a review from antalszava May 25, 2021 15:27

josh146 requested changes May 26, 2021

View reviewed changes

albi3ro changed the title ~~JacobianTape copy duplicates jacobian_options~~ Fixing JAX interface for adjoint diff and finite-diff May 27, 2021

albi3ro mentioned this pull request May 27, 2021

Precision issues when using JAX interface #1363

Closed

albi3ro commented May 27, 2021

View reviewed changes

reset of tape_copy branch

614c29b

albi3ro force-pushed the tape_copy branch from 424228f to 614c29b Compare May 28, 2021 15:53

albi3ro requested a review from josh146 May 28, 2021 15:53

albi3ro and others added 3 commits May 28, 2021 10:55

black

091d7c0

changelog

d7a9c5b

Merge branch 'master' into tape_copy

e853587

antalszava reviewed May 28, 2021

View reviewed changes

albi3ro removed the review-ready 👌 PRs which are ready for review by someone from the core team. label May 28, 2021

josh146 requested changes May 31, 2021

View reviewed changes

remove finite-diff, re-add adjoint test

f8e5ec5

albi3ro changed the title ~~Fixing JAX interface for adjoint diff and finite-diff~~ Fixing JAX interface for adjoint diff May 31, 2021

albi3ro and others added 3 commits May 31, 2021 08:47

Merge remote-tracking branch 'origin/tape_copy' into tape_copy

bbae639

changelog

68734e2

oops

77834fd

another oops

bd4942f

albi3ro requested a review from josh146 May 31, 2021 13:49

restoring accidentally deleted file

e63e3be

josh146 approved these changes May 31, 2021

View reviewed changes

black

97f23b8

albi3ro merged commit a4d7abd into master May 31, 2021

albi3ro deleted the tape_copy branch May 31, 2021 17:14

albi3ro mentioned this pull request May 31, 2021

Adjoint differentiation method fails with jax interface. #1345

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Fixing JAX interface for adjoint diff #1349

Fixing JAX interface for adjoint diff #1349

albi3ro commented May 21, 2021

github-actions bot commented May 21, 2021

codecov bot commented May 21, 2021 •

edited

Loading

albi3ro May 21, 2021

josh146 left a comment

josh146 May 26, 2021

josh146 May 26, 2021

josh146 May 26, 2021

albi3ro May 27, 2021

antalszava May 28, 2021

josh146 May 31, 2021

antalszava May 28, 2021

antalszava commented May 28, 2021

albi3ro commented May 28, 2021

josh146 left a comment •

edited

Loading

josh146 May 31, 2021

albi3ro commented May 31, 2021

josh146 left a comment

josh146 May 31, 2021

josh146 May 31, 2021

		* Fixes drawing QNodes that contain multiple measurements on a single wire.
		[(#1353)](https://github.com/PennyLaneAI/pennylane/pull/1353)

		@@ -55,6 +55,23 @@ def test_parameter_info(self, make_tape):
		}


		class TestTapeCopying:

Fixing JAX interface for adjoint diff #1349

Fixing JAX interface for adjoint diff #1349

Conversation

albi3ro commented May 21, 2021

github-actions bot commented May 21, 2021

codecov bot commented May 21, 2021 • edited Loading

Codecov Report

Choose a reason for hiding this comment

josh146 left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

antalszava commented May 28, 2021

albi3ro commented May 28, 2021

josh146 left a comment • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

albi3ro commented May 31, 2021

josh146 left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

codecov bot commented May 21, 2021 •

edited

Loading

josh146 left a comment •

edited

Loading