Repository: PennyLaneAI/pennylane
Pull request #5620 (2024-05-02 01:45:53)

Make GlobalPhase not differentiable
**Context:**
When using the following state preparation methods (`AmplitudeEmbedding`, `StatePrep`, `MottonenStatePreparation`) with `jit` and `grad`, the error `ValueError: need at least one array to stack` was encountered.

**Description of the Change:**
All state preparation strategies used `GlobalPhase` under the hood, which caused the above error. After this PR, `GlobalPhase` may not be differentiable anymore, as its `grad_method` is set to `None`.

**Benefits:**

**Possible Drawbacks:**

**Related GitHub Issues:**
It fixes  #5541

albi3ro: Thanks for this @Tarun-Kumar07 

For the failures due to errors:

```
ValueError: Computing the gradient of circuits that return the state with the parameter-shift rule gradient transform is not supported, as it is a hardware-compatible method.
```

That would be expected, and we should shift the measurement to expectation values.

For the failures due to:
```
FAILED tests/templates/test_state_preparations/test_mottonen_state_prep.py::test_jacobian_with_and_without_jit_has_same_output_with_high_shots[StatePrep] - AssertionError: assert Array(False, dtype=bool)
 +  where Array(False, dtype=bool) = <function allclose at 0x7fdb77481940>(Array([-0.0003,  0.0153, -0.0153,  0.0003], dtype=float64), Array([ 1.0187, -0.9953, -1.0047,  0.9813], dtype=float64), atol=0.02)
```

Those are legitimately different results, so we can safely safe we are getting wrong results in that case 😢 I'll investigate.

dwierichs: >Those are legitimately different results, so we can safely safe we are getting wrong results in that case 😢 I'll investigate.

@Tarun-Kumar07 @albi3ro Not sure you got to this yet, but it seems that the decomposition of those state preparation methods handle special parameter values differently than others. This makes the derivative wrong at those special values, because `param_shift` is handed a tape that does not contain the general decomposition, so it will not shift all operations that need shifting. As JITting does not allow such special cases, we only make this mistake without JITting, hence the difference in the results within those tests. 

Basically, the decomposition does something like the following decomposition for `RZ`:
```python
def compute_decomposition(theta, wires):
    if not qml.math.is_abstract(theta) and qml.math.isclose(theta, 0):
        return []
    return [qml.RZ(theta, wires)]
```
It's correct but it does not have the correct parameter-shift derivative at 0.

This looks like an independent bug to me, and like one that could be hiding across the codebase for other ops as well, theoretically.
dwierichs: @Tarun-Kumar07 Sorry for taking so long with this! We decided to move ahead with this PR as you originally drafted it (with `grad_method=None`). Before this PR can be merged, the fix in #5774 will need to be merged in, though.
I took the liberty of using a test you wrote here for that PR, your work is much appreciated! :pray: 
Tarun-Kumar07: Hey @dwierichs , once the PR #5774 is merged I will revert changes to `grad_method=None`. Thanks for updating :)
dwierichs: @Tarun-Kumar07 It is merged :) 
I took the liberty to merge the master branch into your branch, updating the tests you wrote (and that made it into #5774) accordingly.
dwierichs: Hi @Tarun-Kumar07 ,
Do you know when you might have time to get back to this? :) Or would you prefer us to finalize it?
Tarun-Kumar07: Hi @dwierichs,

I am currently tied up until August 17th. After that, I will be able to work on this.
If this is a priority, I suggest assigning someone else to take over in the meantime.

Thank you for understanding.


dwierichs: Thanks @Tarun-Kumar07.
That is very understandable of course! I'll try and push your PR through, as it is closed to finish anyways, I think.
Thank you for your contribution, looking forward to the next one :star_struck: 
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5620?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 99.65%. Comparing base [(`9808d0c`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/9808d0ca19f104033970792bbbfe019e236679fc?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`4164ad2`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/4164ad2665a48ef59d8c3952ea996a10eff0d5e3?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).
> Report is 1 commits behind head on master.

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master    #5620      +/-   ##
==========================================
- Coverage   99.66%   99.65%   -0.01%     
==========================================
  Files         430      430              
  Lines       41488    41193     -295     
==========================================
- Hits        41348    41052     -296     
- Misses        140      141       +1     
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5620?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).


Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 58dd4e495..cb77dfa61 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -141,6 +141,11 @@
 
 <h3>Breaking changes 💔</h3>
 
+* `GlobalPhase` is considered non-differentiable with tape transforms.
+  As a consequence, `qml.gradients.finite_diff` and `qml.gradients.spsa_grad` no longer
+  support differentiation of `GlobalPhase` with state-based outputs.
+  [(#5620)](https://github.com/PennyLaneAI/pennylane/pull/5620) 
+
 * The `CircuitGraph.graph` rustworkx graph now stores indices into the circuit as the node labels,
   instead of the operator/ measurement itself.  This allows the same operator to occur multiple times in
   the circuit.
@@ -218,6 +223,8 @@
   [(#5974)](https://github.com/PennyLaneAI/pennylane/pull/5974)
 
 <h3>Bug fixes 🐛</h3>
+* Fix `jax.grad` + `jax.jit` not working for `AmplitudeEmbedding`, `StatePrep` and `MottonenStatePreparation`.
+  [(#5620)](https://github.com/PennyLaneAI/pennylane/pull/5620) 
 
 * Fixed a bug in `qml.SPSAOptimizer` that ignored keyword arguments in the objective function.
   [(#6027)](https://github.com/PennyLaneAI/pennylane/pull/6027)
@@ -246,6 +253,7 @@
 
 This release contains contributions from (in alphabetical order):
 
+Tarun Kumar Allamsetty,
 Guillermo Alonso,
 Utkarsh Azad,
 Ahmed Darwish,
diff --git a/pennylane/ops/identity.py b/pennylane/ops/identity.py
index bdb2d917a..a523c393f 100644
--- a/pennylane/ops/identity.py
+++ b/pennylane/ops/identity.py
@@ -282,7 +282,7 @@ class GlobalPhase(Operation):
 
     """
 
-    grad_method = "A"
+    grad_method = None
     num_params = 1
     num_wires = AllWires
     """int: Number of wires that the operator acts on."""


Test patch:
diff --git a/tests/ops/qubit/test_parametric_ops.py b/tests/ops/qubit/test_parametric_ops.py
index b20b5cc94..304006132 100644
--- a/tests/ops/qubit/test_parametric_ops.py
+++ b/tests/ops/qubit/test_parametric_ops.py
@@ -2015,6 +2015,7 @@ class TestGrad:
         @qml.qnode(dev, diff_method=diff_method)
         def circuit(x):
             qml.Identity(wires[0])
+            qml.GlobalPhase(x, wires=[0, 1])  # Does not change the derivative, but tests it
             qml.Hadamard(wires[1])
             qml.ctrl(qml.GlobalPhase(x), control=wires[1])
             qml.Hadamard(wires[1])
diff --git a/tests/templates/test_embeddings/test_amplitude.py b/tests/templates/test_embeddings/test_amplitude.py
index 859feb114..6dacd4201 100644
--- a/tests/templates/test_embeddings/test_amplitude.py
+++ b/tests/templates/test_embeddings/test_amplitude.py
@@ -584,3 +584,28 @@ class TestInterfaceDtypes:
         res2 = circuit2(features, pad_with, normalize=True)
 
         assert qml.math.allclose(res, res2, atol=tol, rtol=0)
+
+
+@pytest.mark.jax
+@pytest.mark.parametrize("shots, atol", [(10000, 0.05), (None, 1e-8)])
+def test_jacobian_with_and_without_jit_has_same_output(shots, atol):
+    """Test that the jacobian of AmplitudeEmbedding is the same with and without jit."""
+
+    import jax
+
+    dev = qml.device("default.qubit", shots=shots, seed=7890234)
+
+    @qml.qnode(dev, diff_method="parameter-shift")
+    def circuit(coeffs):
+        qml.AmplitudeEmbedding(coeffs, normalize=True, wires=[0, 1])
+        return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))
+
+    params = jax.numpy.array([0.4, 0.5, 0.1, 0.3])
+    jac_fn = jax.jacobian(circuit)
+    jac_jit_fn = jax.jit(jac_fn)
+
+    jac = jac_fn(params)
+
+    jac_jit = jac_jit_fn(params)
+
+    assert qml.math.allclose(jac, jac_jit, atol=atol)
diff --git a/tests/templates/test_state_preparations/test_mottonen_state_prep.py b/tests/templates/test_state_preparations/test_mottonen_state_prep.py
index d8c4cdaf2..d1c0e194c 100644
--- a/tests/templates/test_state_preparations/test_mottonen_state_prep.py
+++ b/tests/templates/test_state_preparations/test_mottonen_state_prep.py
@@ -430,16 +430,21 @@ def test_jacobians_with_and_without_jit_match(shots, atol):
         return qml.probs(wires=[0, 1])
 
     circuit_fd = qml.QNode(circuit, dev, diff_method="finite-diff", h=0.05)
+    circuit_ps = qml.QNode(circuit, dev, diff_method="parameter-shift")
     circuit_exact = qml.QNode(circuit, dev_no_shots)
 
     params = jax.numpy.array([0.5, 0.5, 0.5, 0.5])
     jac_exact_fn = jax.jacobian(circuit_exact)
-    jac_fn = jax.jacobian(circuit_fd)
-    jac_jit_fn = jax.jit(jac_fn)
+    jac_fd_fn = jax.jacobian(circuit_fd)
+    jac_fd_fn_jit = jax.jit(jac_fd_fn)
+    jac_ps_fn = jax.jacobian(circuit_ps)
+    jac_ps_fn_jit = jax.jit(jac_ps_fn)
 
     jac_exact = jac_exact_fn(params)
-    jac = jac_fn(params)
-    jac_jit = jac_jit_fn(params)
+    jac_fd = jac_fd_fn(params)
+    jac_fd_jit = jac_fd_fn_jit(params)
+    jac_ps = jac_ps_fn(params)
+    jac_ps_jit = jac_ps_fn_jit(params)
 
-    assert qml.math.allclose(jac_exact, jac_jit, atol=atol)
-    assert qml.math.allclose(jac, jac_jit, atol=atol)
+    for compare in [jac_fd, jac_fd_jit, jac_ps, jac_ps_jit]:
+        assert qml.math.allclose(jac_exact, compare, atol=atol)
diff --git a/tests/transforms/test_tape_expand.py b/tests/transforms/test_tape_expand.py
index 4bf16a061..3d57d9207 100644
--- a/tests/transforms/test_tape_expand.py
+++ b/tests/transforms/test_tape_expand.py
@@ -336,9 +336,8 @@ class TestExpandInvalidTrainable:
 
         assert new_tape.operations[0].name == "RZ"
         assert new_tape.operations[0].grad_method == "A"
-        assert new_tape.operations[1].name == "GlobalPhase"
-        assert new_tape.operations[2].name == "RY"
-        assert new_tape.operations[3].name == "CNOT"
+        assert new_tape.operations[1].name == "RY"
+        assert new_tape.operations[2].name == "CNOT"
 
     def test_nontrainable_nondiff(self, mocker):
         """Test that a circuit with non-differentiable