Repository: PennyLaneAI/pennylane
Pull request #1349 (2021-05-21 14:28:14)

Fixing JAX interface for adjoint diff
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[bot]: Hello. You may have forgotten to update the changelog!
Please edit [.github/CHANGELOG.md](/PennyLaneAI/pennylane/blob/master/.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.
codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1349?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#1349](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1349?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (77d08d1) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/7e8ddfb1b3203c9aa21a29c924f890ad7c8c9f3f?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (7e8ddfb) will **increase** coverage by `0.00%`.
> The diff coverage is `100.00%`.

[![Impacted file tree graph](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1349/graphs/tree.svg?width=650&height=150&src=pr&token=9S3YzlUgGk&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1349?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)

```diff
@@           Coverage Diff           @@
##           master    #1349   +/-   ##
=======================================
  Coverage   98.25%   98.26%           
=======================================
  Files         146      146           
  Lines       11146    11150    +4     
=======================================
+ Hits        10952    10956    +4     
  Misses        194      194           
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1349?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/tape/jacobian\_tape.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1349/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RhcGUvamFjb2JpYW5fdGFwZS5weQ==) | `97.94% <100.00%> (+0.04%)` | :arrow_up: |
| [pennylane/ops/qubit.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1349/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL29wcy9xdWJpdC5weQ==) | `98.56% <0.00%> (ø)` | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1349?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1349?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Last update [7e8ddfb...77d08d1](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1349?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).

albi3ro: 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.
albi3ro: 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.

Patch:
diff --git a/.github/CHANGELOG.md b/.github/CHANGELOG.md
index 77963e570..3fd197299 100644
--- a/.github/CHANGELOG.md
+++ b/.github/CHANGELOG.md
@@ -345,6 +345,10 @@ random_mat2 = rng.standard_normal(3, requires_grad=False)
 
 <h3>Bug fixes</h3>
 
+* Copying the `JacobianTape` now correctly also copies the `jacobian_options` attribute. This fixes a bug
+  allowing the JAX interface to support adjoint differentiation.
+  [(#1349)](https://github.com/PennyLaneAI/pennylane/pull/1349)
+
 * Fixes drawing QNodes that contain multiple measurements on a single wire.
   [(#1353)](https://github.com/PennyLaneAI/pennylane/pull/1353)
 
@@ -393,7 +397,7 @@ random_mat2 = rng.standard_normal(3, requires_grad=False)
 This release contains contributions from (in alphabetical order):
 
 Marius Aglitoiu, Vishnu Ajith, Thomas Bromley, Jack Ceroni, Miruna Daian, Olivia Di Matteo,
-Tanya Garg, Christian Gogolin, Diego Guala, Anthony Hayes, Ryan Hill, Josh Izaac, Pavan Jayasinha, Ryan Levy, Nahum Sá, Maria Schuld,
+Tanya Garg, Christian Gogolin, Diego Guala, Anthony Hayes, Ryan Hill, Josh Izaac, Pavan Jayasinha, Christina Lee, Ryan Levy, Nahum Sá, Maria Schuld,
 Johannes Jakob Meyer, Brian Shi, Antal Száva, David Wierichs, Vincent Wong, Alberto Maldonado, Ashish Panigrahi.
 
 
diff --git a/pennylane/interfaces/autograd.py b/pennylane/interfaces/autograd.py
index ec138282f..80211acc5 100644
--- a/pennylane/interfaces/autograd.py
+++ b/pennylane/interfaces/autograd.py
@@ -119,7 +119,7 @@ class AutogradInterface(AnnotatedQueue):
                 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]))
 
         By default, all parameters are trainable and will be returned:
 
diff --git a/pennylane/tape/jacobian_tape.py b/pennylane/tape/jacobian_tape.py
index fee1b3255..9296ea175 100644
--- a/pennylane/tape/jacobian_tape.py
+++ b/pennylane/tape/jacobian_tape.py
@@ -93,6 +93,11 @@ class JacobianTape(QuantumTape):
         self.jacobian_options = {}
         self.hessian_options = {}
 
+    def copy(self, copy_operations=False, tape_cls=None):
+        copied_tape = super().copy(copy_operations=copy_operations, tape_cls=tape_cls)
+        copied_tape.jacobian_options = self.jacobian_options
+        return copied_tape
+
     def _grad_method(self, idx, use_graph=True, default_method="F"):
         """Determine the correct partial derivative computation method for each gate parameter.
 


Test patch:
diff --git a/tests/tape/test_jacobian_tape.py b/tests/tape/test_jacobian_tape.py
index 7ab75effc..fd9aa7a90 100644
--- a/tests/tape/test_jacobian_tape.py
+++ b/tests/tape/test_jacobian_tape.py
@@ -55,6 +55,23 @@ class TestConstruction:
         }
 
 
+class TestTapeCopying:
+    """Test for tape copying behaviour"""
+
+    def test_jacobian_options_copied(self):
+        """Tests that the jacobian_options attribute is copied"""
+
+        tape = JacobianTape()
+        tape.jacobian_options = {"method": "device", "jacobian_method": "adjoint_jacobian"}
+
+        tape_copy = tape.copy()
+
+        assert tape_copy.jacobian_options == {
+            "method": "device",
+            "jacobian_method": "adjoint_jacobian",
+        }
+
+
 class TestGradMethod:
     """Tests for parameter gradient methods"""
 
diff --git a/tests/test_qubit_device_adjoint_jacobian.py b/tests/test_qubit_device_adjoint_jacobian.py
index 8f4226e8d..9f205084c 100644
--- a/tests/test_qubit_device_adjoint_jacobian.py
+++ b/tests/test_qubit_device_adjoint_jacobian.py
@@ -374,3 +374,25 @@ class TestAdjointJacobianQNode:
         grad_fd = params1.grad, params2.grad
 
         assert np.allclose(grad_adjoint, grad_fd)
+
+    def test_interface_jax(self, dev):
+        """Test if the gradients agree between adjoint and backprop methods in the
+        jax interface"""
+        jax = pytest.importorskip("jax")
+
+        def f(params1, params2):
+            qml.RX(0.4, wires=[0])
+            qml.RZ(params1 * jax.numpy.sqrt(params2), wires=[0])
+            qml.RY(jax.numpy.cos(params2), wires=[0])
+            return qml.expval(qml.PauliZ(0))
+
+        params1 = jax.numpy.array(0.3)
+        params2 = jax.numpy.array(0.4)
+
+        qnode_adjoint = QNode(f, dev, interface="jax", diff_method="adjoint")
+        qnode_backprop = QNode(f, dev, interface="jax", diff_method="backprop")
+
+        grad_adjoint = jax.grad(qnode_adjoint)(params1, params2)
+        grad_backprop = jax.grad(qnode_backprop)(params1, params2)
+
+        assert np.allclose(grad_adjoint, grad_backprop)