Repository: PennyLaneAI/pennylane
Pull request #2289 (2022-03-07 18:28:42)

[Bug] Exclude Snapshot from adjoint backwards pass
The adjoint diff-method on `QubitDevice`s bypasses the normal `dev.apply` function and instead invokes `dev._apply_operation` for all operations on the tape. This leads to a bug with the `Snapshot` operation as implemented on the `default.qubit` device.
See the following example:

```py
dev = qml.device("default.qubit", wires=1)

@qml.qnode(dev, diff_method="adjoint")
def circuit():
    qml.Hadamard(0)
    qml.Snapshot()
    return qml.expval(qml.PauliZ(0))

circuit()
```
```
ValueError: cannot reshape array of size 4 into shape ()
```

github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [`doc/releases/changelog-dev.md`](/PennyLaneAI/pennylane/blob/master/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[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2289?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#2289](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2289?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (dcd7933) into [v0.22.0-rc0](https://codecov.io/gh/PennyLaneAI/pennylane/commit/81b9c1f6db4c0ae750d93d4554417616e186ed3c?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (81b9c1f) will **decrease** coverage by `2.61%`.
> The diff coverage is `100.00%`.

[![Impacted file tree graph](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2289/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/2289?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)

```diff
@@               Coverage Diff               @@
##           v0.22.0-rc0    #2289      +/-   ##
===============================================
- Coverage        99.31%   96.70%   -2.62%     
===============================================
  Files              241      241              
  Lines            19101    19101              
===============================================
- Hits             18971    18471     -500     
- Misses             130      630     +500     
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2289?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/\_qubit\_device.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2289/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL19xdWJpdF9kZXZpY2UucHk=) | `98.40% <100.00%> (-0.32%)` | :arrow_down: |
| [pennylane/interfaces/batch/jax.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2289/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvYmF0Y2gvamF4LnB5) | `0.00% <0.00%> (-100.00%)` | :arrow_down: |
| [pennylane/interfaces/batch/jax\_jit.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2289/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvYmF0Y2gvamF4X2ppdC5weQ==) | `0.00% <0.00%> (-100.00%)` | :arrow_down: |
| [pennylane/interfaces/jax.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2289/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvamF4LnB5) | `0.00% <0.00%> (-97.44%)` | :arrow_down: |
| [pennylane/devices/default\_qubit\_jax.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2289/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RldmljZXMvZGVmYXVsdF9xdWJpdF9qYXgucHk=) | `0.00% <0.00%> (-96.63%)` | :arrow_down: |
| [pennylane/qnn/cost.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2289/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3Fubi9jb3N0LnB5) | `42.85% <0.00%> (-57.15%)` | :arrow_down: |
| [pennylane/transforms/classical\_jacobian.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2289/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvY2xhc3NpY2FsX2phY29iaWFuLnB5) | `77.08% <0.00%> (-22.92%)` | :arrow_down: |
| [pennylane/math/utils.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2289/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21hdGgvdXRpbHMucHk=) | `84.41% <0.00%> (-15.59%)` | :arrow_down: |
| [pennylane/math/multi\_dispatch.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2289/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21hdGgvbXVsdGlfZGlzcGF0Y2gucHk=) | `87.73% <0.00%> (-12.27%)` | :arrow_down: |
| [pennylane/math/is\_independent.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2289/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21hdGgvaXNfaW5kZXBlbmRlbnQucHk=) | `88.50% <0.00%> (-11.50%)` | :arrow_down: |
| ... and [14 more](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2289/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2289?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/2289?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Last update [81b9c1f...dcd7933](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2289?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).

dime10: @antalszava CodeFactor seems hung up. I ran pylint locally, no new issues.

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 177b16afc..7f3f2a216 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -9,7 +9,8 @@
   execution, such as the quantum state vector and density matrix in the qubit case, or the
   covariance matrix and vector of means in the continuous variable case. The saved states can be
   retrieved in the form of a dictionary via the top-level `qml.snapshots` function.
-  [(#2233)](https://github.com/PennyLaneAI/pennylane/pull/2233)
+  [(#2233)](https://github.com/PennyLaneAI/pennylane/pull/2233) 
+  [(#2289)](https://github.com/PennyLaneAI/pennylane/pull/2289)
 
   ```py
   dev = qml.device("default.qubit", wires=2)
diff --git a/pennylane/_qubit_device.py b/pennylane/_qubit_device.py
index 450f56210..1c5d6b73b 100644
--- a/pennylane/_qubit_device.py
+++ b/pennylane/_qubit_device.py
@@ -942,7 +942,7 @@ class QubitDevice(Device):
                         'the "adjoint" differentiation method'
                     )
             else:
-                if op.name not in ("QubitStateVector", "BasisState"):
+                if op.name not in ("QubitStateVector", "BasisState", "Snapshot"):
                     expanded_ops.append(op)
 
         jac = np.zeros((len(tape.observables), len(tape.trainable_params)))


Test patch:
diff --git a/tests/test_debugging.py b/tests/test_debugging.py
index 69ea53993..1a933f3f8 100644
--- a/tests/test_debugging.py
+++ b/tests/test_debugging.py
@@ -252,3 +252,28 @@ class TestSnapshot:
 
         if m == "state":
             assert np.allclose(result[2], result["execution_results"])
+
+    @pytest.mark.parametrize("method", [None, "backprop", "parameter-shift", "adjoint"])
+    def test_different_diff_methods(self, method):
+        """Test that snapshots work with different differentiation methods."""
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev, diff_method=method)
+        def circuit():
+            qml.Snapshot()
+            qml.Hadamard(wires=0)
+            qml.Snapshot("very_important_state")
+            qml.CNOT(wires=[0, 1])
+            qml.Snapshot()
+            return qml.expval(qml.PauliZ(0))
+
+        result = qml.snapshots(circuit)()
+        expected = {
+            0: np.array([1, 0, 0, 0]),
+            "very_important_state": np.array([1 / np.sqrt(2), 0, 1 / np.sqrt(2), 0]),
+            2: np.array([1 / np.sqrt(2), 0, 0, 1 / np.sqrt(2)]),
+            "execution_results": np.array(0),
+        }
+
+        assert all(k1 == k2 for k1, k2 in zip(result.keys(), expected.keys()))
+        assert all(np.allclose(v1, v2) for v1, v2 in zip(result.values(), expected.values()))