Repository: PennyLaneAI/pennylane
Pull request #4706 (2023-10-23 15:11:25)

Hybrid kwargs for gradient transform
**Description of the Change:**

Readd hybrid kwarg for backward compatibility of gradient transforms.


codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4706?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 [(`aa626e3`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/aa626e3bdf08105faed9ce5299e308096fb478c0?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.64% compared to head [(`ba85ac2`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4706?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.63%.

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


```diff
@@            Coverage Diff             @@
##           master    #4706      +/-   ##
==========================================
- Coverage   99.64%   99.63%   -0.01%     
==========================================
  Files         378      378              
  Lines       34071    33812     -259     
==========================================
- Hits        33950    33690     -260     
- Misses        121      122       +1     
```


| [Files](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4706?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/qnode.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4706?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3Fub2RlLnB5) | `100.00% <100.00%> (ø)` | |
| [pennylane/transforms/core/transform\_program.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4706?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvY29yZS90cmFuc2Zvcm1fcHJvZ3JhbS5weQ==) | `99.46% <100.00%> (ø)` | |

... and [42 files with indirect coverage changes](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4706/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)


</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4706?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/pennylane/qnode.py b/pennylane/qnode.py
index 90555e3d9..1a9362f14 100644
--- a/pennylane/qnode.py
+++ b/pennylane/qnode.py
@@ -998,16 +998,22 @@ class QNode:
             )
         # Calculate the classical jacobians if necessary
         if full_transform_program.has_classical_cotransform():
-            argnums = full_transform_program[-1]._kwargs.pop(
-                "argnums", None
-            )  # pylint: disable=protected-access
-            full_transform_program._set_all_classical_jacobians(
-                self, args, kwargs, argnums
-            )  # pylint: disable=protected-access
-            full_transform_program._set_all_argnums(
-                self, args, kwargs, argnums
+            hybrid = full_transform_program[-1]._kwargs.pop(
+                "hybrid", True
             )  # pylint: disable=protected-access
 
+            if hybrid:
+                argnums = full_transform_program[-1]._kwargs.pop(
+                    "argnums", None
+                )  # pylint: disable=protected-access
+
+                full_transform_program._set_all_classical_jacobians(
+                    self, args, kwargs, argnums
+                )  # pylint: disable=protected-access
+                full_transform_program._set_all_argnums(
+                    self, args, kwargs, argnums
+                )  # pylint: disable=protected-access
+
         # pylint: disable=unexpected-keyword-arg
         res = qml.execute(
             (self._tape,),
diff --git a/pennylane/transforms/core/transform_program.py b/pennylane/transforms/core/transform_program.py
index 21f4ab2b6..d6ee51fc9 100644
--- a/pennylane/transforms/core/transform_program.py
+++ b/pennylane/transforms/core/transform_program.py
@@ -434,7 +434,7 @@ class TransformProgram:
                     slices_classical.append(slice(start_classical, start_classical + 1))
                     start_classical += 1
 
-            if cotransform:
+            if cotransform and self._classical_jacobians:
                 batch_postprocessing_classical = partial(
                     _batch_postprocessing, individual_fns=classical_fns, slices=slices_classical
                 )


Test patch:
diff --git a/tests/gradients/parameter_shift/test_parameter_shift.py b/tests/gradients/parameter_shift/test_parameter_shift.py
index d4feab4d4..829cc776a 100644
--- a/tests/gradients/parameter_shift/test_parameter_shift.py
+++ b/tests/gradients/parameter_shift/test_parameter_shift.py
@@ -3464,6 +3464,25 @@ class TestQnodeAutograd:
         assert res.shape == res_expected.shape
         assert np.allclose(res, res_expected)
 
+    @pytest.mark.parametrize("interface", interfaces)
+    def test_single_measurement_single_param_not_hybrid(self, interface):
+        """Test for a single measurement and a single param with hybrid False."""
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev, interface=interface)
+        def circuit(x):
+            qml.RX(2 * x, wires=[0])
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0))
+
+        x = qml.numpy.array(0.543, requires_grad=True)
+
+        res = qml.gradients.param_shift(circuit, hybrid=False)(x)
+
+        res_expected = qml.jacobian(circuit)(x)
+        assert res.shape == res_expected.shape
+        assert np.allclose(2 * res, res_expected)
+
     @pytest.mark.parametrize("interface", interfaces)
     def test_single_measurement_single_param_2(self, interface):
         """Test for a single measurement and a single param."""