Repository: PennyLaneAI/pennylane
Pull request #4879 (2023-11-24 17:41:44)

Fix number trainable params in finite diff
**Context:**

When using finite diff directly on a QNode with a diff method and with Jax might results in switching the argnums between the transform and its post processing.

**Description of the Change:**

Set the trainable parameters in the transforms and use then in the post processing.


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://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4879?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:
> Comparison is base [(`370d9d4`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/370d9d447bf47b6060369ed3d36f00ba1097c82f?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.65% compared to head [(`019e570`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4879?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.64%.


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


```diff
@@            Coverage Diff             @@
##           master    #4879      +/-   ##
==========================================
- Coverage   99.65%   99.64%   -0.02%     
==========================================
  Files         385      385              
  Lines       34719    34461     -258     
==========================================
- Hits        34599    34338     -261     
- Misses        120      123       +3     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4879?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).

rmoyard: When setting diff_method="finite_diff" it chooses the execution with the Jax interface. In the Jax interface we set the trainable parameters, as None of them are considered trainable because there are not tracers `isinstance(tensor, jax.core.Tracer)`. Indeed nothing is jacobian or grad is not used. Users should not choose a diff method when applying a transform but that's possible.

There is a stability issue with Torch, the execution pipeline default vs torch have different results.
trbromley: > Also, does this PR have anything to do with Issue https://github.com/PennyLaneAI/pennylane/issues/4882 ?

They're related with similar example code (except the interface). In #4882, Torch gives incorrect gradients when `diff_method` is specified. Here, the issue is that JAX raises an error when `diff_method` is _not_ specified. :confused:

TLDR I don't think this PR helps with #4882.

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index f2d034d36..ee46783b6 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -295,6 +295,10 @@
 
 <h3>Bug fixes 🐛</h3>
 
+* Fix a bug where trainable parameters in the post-processing of finite diff were incorrect for Jax when applying
+  the transform directly on a ``QNode``.
+  [(#4879)](https://github.com/PennyLaneAI/pennylane/pull/4879)
+
 * `qml.grad` and `qml.jacobian` now explicitly raise errors if trainable parameters are integers.
   [(#4836)](https://github.com/PennyLaneAI/pennylane/pull/4836)
 
@@ -381,6 +385,7 @@ Ankit Khandelwal,
 Christina Lee,
 Romain Moyard,
 Vincent Michaud-Rioux,
+Romain Moyard,
 Anurav Modak,
 Mudit Pandey,
 Matthew Silverman,
diff --git a/pennylane/gradients/finite_difference.py b/pennylane/gradients/finite_difference.py
index 767a2ce66..66bb2f562 100644
--- a/pennylane/gradients/finite_difference.py
+++ b/pennylane/gradients/finite_difference.py
@@ -359,7 +359,8 @@ def finite_diff(
             "Finite differences with float32 detected. Answers may be inaccurate. float64 is recommended.",
             UserWarning,
         )
-
+    number_parameters = len(tape.trainable_params)
+    number_measurements = len(tape.measurements)
     if argnum is None and not tape.trainable_params:
         return _no_trainable_grad(tape)
 
@@ -408,7 +409,6 @@ def finite_diff(
     def _single_shot_batch_result(results):
         """Auxiliary function for post-processing one batch of results corresponding to finite shots or a single
         component of a shot vector"""
-
         grads = []
         start = 1 if c0 is not None and f0 is None else 0
         r0 = f0 or results[0]
@@ -443,13 +443,13 @@ def finite_diff(
 
             pre_grads = []
 
-            if len(tape.measurements) == 1:
+            if number_measurements == 1:
                 res = qml.math.stack(res)
                 c = qml.math.convert_like(coeffs, res)
                 lin_comb = qml.math.tensordot(res, c, [[0], [0]])
                 pre_grads.append(lin_comb)
             else:
-                for i in range(len(tape.measurements)):
+                for i in range(number_measurements):
                     r = qml.math.stack([r[i] for r in res])
                     c = qml.math.convert_like(coeffs, r)
                     lin_comb = qml.math.tensordot(r, c, [[0], [0]])
@@ -457,11 +457,11 @@ def finite_diff(
 
             # Add on the unshifted term
             if c0 is not None:
-                if len(tape.measurements) == 1:
+                if number_measurements == 1:
                     c = qml.math.convert_like(c0, r0)
                     pre_grads = [pre_grads[0] + r0 * c]
                 else:
-                    for i in range(len(tape.measurements)):
+                    for i in range(number_measurements):
                         r_i = r0[i]
                         c = qml.math.convert_like(c0, r_i)
                         pre_grads[i] = pre_grads[i] + r_i * c
@@ -479,19 +479,19 @@ def finite_diff(
 
             grads.append(pre_grads)
         # Single measurement
-        if len(tape.measurements) == 1:
-            if len(tape.trainable_params) == 1:
+        if number_measurements == 1:
+            if number_parameters == 1:
                 return grads[0]
             return tuple(grads)
 
         # Reordering to match the right shape for multiple measurements
-        grads_reorder = [[0] * len(tape.trainable_params) for _ in range(len(tape.measurements))]
-        for i in range(len(tape.measurements)):
-            for j in range(len(tape.trainable_params)):
+        grads_reorder = [[0] * number_parameters for _ in range(len(tape.measurements))]
+        for i in range(number_measurements):
+            for j in range(number_parameters):
                 grads_reorder[i][j] = grads[j][i]
 
         # To tuple
-        if len(tape.trainable_params) == 1:
+        if number_parameters == 1:
             return tuple(elem[0] for elem in grads_reorder)
         return tuple(tuple(elem) for elem in grads_reorder)
 


Test patch:
diff --git a/tests/gradients/finite_diff/test_finite_difference.py b/tests/gradients/finite_diff/test_finite_difference.py
index 9f1752dba..6a5f71082 100644
--- a/tests/gradients/finite_diff/test_finite_difference.py
+++ b/tests/gradients/finite_diff/test_finite_difference.py
@@ -1065,6 +1065,46 @@ class TestFiniteDiffGradients:
         )
         assert np.allclose(res, expected, atol=tol, rtol=0)
 
+    @pytest.mark.jax
+    @pytest.mark.parametrize("dev_name", ["default.qubit"])
+    def test_jax_probs(
+        self, dev_name, approx_order, strategy, tol
+    ):  # pylint: disable=unused-argument
+        """Tests that the output of the finite-difference transform is similar using or not diff method on the QNode."""
+        import jax
+        from jax.config import config
+
+        config.update("jax_enable_x64", True)
+
+        dev = qml.device(dev_name, wires=2)
+        x = jax.numpy.array(0.543)
+
+        @qml.qnode(dev)
+        def circuit(x):
+            qml.RX(x, wires=[0])
+            qml.RY(x, wires=[0])
+            qml.CNOT(wires=[0, 1])
+            return qml.probs(wires=0), qml.probs(wires=1)
+
+        @qml.qnode(dev, diff_method="finite-diff")
+        def circuit_fd(x):
+            qml.RX(x, wires=[0])
+            qml.RY(x, wires=[0])
+            qml.CNOT(wires=[0, 1])
+            return qml.probs(wires=0), qml.probs(wires=1)
+
+        transform_output = qml.gradients.finite_diff(circuit)(x)
+        framework_output = jax.jacobian(circuit)(x)
+        assert jax.numpy.allclose(
+            jax.numpy.stack(transform_output), jax.numpy.stack(framework_output)
+        )
+
+        transform_output = qml.gradients.finite_diff(circuit_fd)(x)
+        framework_output = jax.jacobian(circuit_fd)(x)
+        assert jax.numpy.allclose(
+            jax.numpy.stack(transform_output), jax.numpy.stack(framework_output)
+        )
+
 
 @pytest.mark.parametrize("argnums", [[0], [1], [0, 1]])
 @pytest.mark.parametrize("interface", ["jax"])
diff --git a/tests/gradients/finite_diff/test_finite_difference_shot_vec.py b/tests/gradients/finite_diff/test_finite_difference_shot_vec.py
index 039d280fb..68c17d0e6 100644
--- a/tests/gradients/finite_diff/test_finite_difference_shot_vec.py
+++ b/tests/gradients/finite_diff/test_finite_difference_shot_vec.py
@@ -774,7 +774,7 @@ class TestFiniteDiffIntegration:
 
             assert isinstance(res[0], tuple)
             assert len(res[0]) == 2
-            assert np.allclose(res[0], [-np.sin(x), 0], atol=0.1, rtol=0)
+            assert np.allclose(res[0], [-np.sin(x), 0], atol=0.2, rtol=0)
             assert isinstance(res[0][0], numpy.ndarray)
             assert isinstance(res[0][1], numpy.ndarray)