Repository: PennyLaneAI/pennylane
Pull request #6898 (2025-01-29 14:43:32)

Remove internal `QNode.get_gradient_fn` usage
**Context:**

`QNode.get_gradient_fn` was due for deprecation last release but at the last minute it was halted as there was no alternative to propose to users. Now with the workflow clean-up, we have functions that can be used as an alternative (at least internally).

**Description of the Change:**

- Added a `_validate_diff_method` helper function to `qnode.py`
- Replaced usage of `QNode.get_gradient_fn` with `qml.workflow.resolution._resolve_diff_method` where appropriate

After replacing usage, we lost project coverage 🤣 .

- Added `TestGetGradientFn` test class with additional tests

**Benefits:**

Prepares `QNode.get_gradient_fn` for removal.

**Possible Drawbacks:** None.

[sc-83315]


codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/6898?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.56%. Comparing base [(`06cd2f0`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/06cd2f0d01ff67495acc566ff0244202da7e48c8?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`00c73f1`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/00c73f151a9fa3d3a6d845ac4d213e17d02a4439?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).
> Report is 2 commits behind head on master.

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


```diff
@@            Coverage Diff             @@
##           master    #6898      +/-   ##
==========================================
- Coverage   99.59%   99.56%   -0.04%     
==========================================
  Files         478      478              
  Lines       45306    45324      +18     
==========================================
+ Hits        45124    45128       +4     
- Misses        182      196      +14     
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/6898?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 a66b7802e..cba18c600 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -181,7 +181,10 @@
 
 <h3>Internal changes ⚙️</h3>
 
-* The source code has been updated use black 25.1.0
+* Remove `QNode.get_gradient_fn` from source code.
+  [(#6898)](https://github.com/PennyLaneAI/pennylane/pull/6898)
+  
+* The source code has been updated use black 25.1.0.
   [(#6897)](https://github.com/PennyLaneAI/pennylane/pull/6897)
 
 * Improved the `InterfaceEnum` object to prevent direct comparisons to `str` objects.
diff --git a/pennylane/resource/specs.py b/pennylane/resource/specs.py
index c83f84c7a..b14874f52 100644
--- a/pennylane/resource/specs.py
+++ b/pennylane/resource/specs.py
@@ -214,12 +214,8 @@ def specs(
                 else qnode.diff_method
             )
 
-            gradient_fn = qml.QNode.get_gradient_fn(
-                qnode.device,
-                qnode.interface,
-                qnode.diff_method,
-                tape=tape,
-            )[0]
+            config = qml.workflow.construct_execution_config(qnode)(*args, **kwargs)
+            gradient_fn = config.gradient_method
             if isinstance(gradient_fn, qml.transforms.core.TransformDispatcher):
                 info["gradient_fn"] = _get_absolute_import_path(gradient_fn)
 
diff --git a/pennylane/workflow/construct_batch.py b/pennylane/workflow/construct_batch.py
index 8451b8d77..f74fba117 100644
--- a/pennylane/workflow/construct_batch.py
+++ b/pennylane/workflow/construct_batch.py
@@ -184,7 +184,13 @@ def get_transform_program(
 
     """
     if gradient_fn == "unset":
-        gradient_fn = QNode.get_gradient_fn(qnode.device, qnode.interface, qnode.diff_method)[0]
+        config = qml.workflow.construct_execution_config(qnode, resolve=False)()
+        # pylint: disable = protected-access
+        config = qml.workflow.resolution._resolve_diff_method(
+            config,
+            qnode.device,
+        )
+        gradient_fn = config.gradient_method
 
     full_transform_program = _get_full_transform_program(qnode, gradient_fn)
 
@@ -362,9 +368,12 @@ def construct_batch(
             params = initial_tape.get_parameters(trainable_only=False)
             initial_tape.trainable_params = qml.math.get_trainable_indices(params)
 
-        gradient_fn = QNode.get_gradient_fn(
-            qnode.device, qnode.interface, qnode.diff_method, tape=initial_tape
-        )[0]
+        config = qml.workflow.construct_execution_config(qnode, resolve=False)(*args, **kwargs)
+        # pylint: disable = protected-access
+        config = qml.workflow.resolution._resolve_execution_config(
+            config, qnode.device, tapes=(initial_tape,)
+        )
+        gradient_fn = config.gradient_method
         program = get_transform_program(qnode, level=level, gradient_fn=gradient_fn)
 
         return program((initial_tape,))
diff --git a/pennylane/workflow/construct_execution_config.py b/pennylane/workflow/construct_execution_config.py
index 0b13700d5..1dfba42f7 100644
--- a/pennylane/workflow/construct_execution_config.py
+++ b/pennylane/workflow/construct_execution_config.py
@@ -81,7 +81,6 @@ def construct_execution_config(qnode: "qml.QNode", resolve: bool = True):
 
     @functools.wraps(qnode)
     def wrapper(*args, **kwargs):
-
         mcm_config = qml.devices.MCMConfig(
             postselect_mode=qnode.execute_kwargs["postselect_mode"],
             mcm_method=qnode.execute_kwargs["mcm_method"],
@@ -104,7 +103,7 @@ def construct_execution_config(qnode: "qml.QNode", resolve: bool = True):
         )
 
         if resolve:
-            tape = construct_tape(qnode)(*args, **kwargs)
+            tape = construct_tape(qnode, level=0)(*args, **kwargs)
             # pylint:disable=protected-access
             config = _resolve_execution_config(
                 config, qnode.device, (tape,), qnode._transform_program
diff --git a/pennylane/workflow/qnode.py b/pennylane/workflow/qnode.py
index 4d2d6789e..594d54dff 100644
--- a/pennylane/workflow/qnode.py
+++ b/pennylane/workflow/qnode.py
@@ -178,6 +178,32 @@ def _validate_qfunc_output(qfunc_output, measurements) -> None:
         )
 
 
+def _validate_diff_method(
+    device: SupportedDeviceAPIs, diff_method: Union[str, TransformDispatcher]
+) -> None:
+    if diff_method is None:
+        return
+
+    # performs type validation
+    config = _make_execution_config(None, diff_method)
+
+    if device.supports_derivatives(config):
+        return
+    if diff_method in {"backprop", "adjoint", "device"}:  # device-only derivatives
+        raise qml.QuantumFunctionError(
+            f"Device {device} does not support {diff_method} with requested circuit."
+        )
+    if isinstance(diff_method, str) and diff_method in tuple(get_args(SupportedDiffMethods)):
+        return
+    if isinstance(diff_method, TransformDispatcher):
+        return
+
+    raise qml.QuantumFunctionError(
+        f"Differentiation method {diff_method} not recognized. Allowed "
+        f"options are {tuple(get_args(SupportedDiffMethods))}."
+    )
+
+
 # pylint: disable=too-many-instance-attributes
 class QNode:
     r"""Represents a quantum node in the hybrid computational graph.
@@ -572,6 +598,7 @@ class QNode:
         self.device = device
         self._interface = get_canonical_interface_name(interface)
         self.diff_method = diff_method
+        _validate_diff_method(self.device, self.diff_method)
         cache = (max_diff > 1) if cache == "auto" else cache
 
         # execution keyword arguments
@@ -595,10 +622,6 @@ class QNode:
         self._transform_program = TransformProgram()
         functools.update_wrapper(self, func)
 
-        # validation check.  Will raise error if bad diff_method
-        if diff_method is not None:
-            QNode.get_gradient_fn(self.device, self.interface, self.diff_method)
-
     def __copy__(self) -> "QNode":
         copied_qnode = QNode.__new__(QNode)
         for attr, value in vars(self).items():


Test patch:
diff --git a/tests/logging/test_logging_autograd.py b/tests/logging/test_logging_autograd.py
index 47ae788f3..48861f41e 100644
--- a/tests/logging/test_logging_autograd.py
+++ b/tests/logging/test_logging_autograd.py
@@ -106,7 +106,7 @@ class TestLogging:
                 assert msg in actual.getMessage()
 
     @pytest.mark.parametrize(
-        "diff_method,num_records", [("parameter-shift", 24), ("backprop", 15), ("adjoint", 19)]
+        "diff_method,num_records", [("parameter-shift", 23), ("backprop", 13), ("adjoint", 17)]
     )
     def test_dq_qnode_execution_grad(self, caplog, diff_method, num_records):
         "Test logging of QNode with parametrized gradients"
@@ -135,12 +135,6 @@ class TestLogging:
                     f"interface=Interface.AUTO, diff_method={diff_method}, grad_on_execution=best,",
                 ],
             ),
-            (
-                "pennylane.workflow.qnode",
-                [
-                    "Calling <get_gradient_fn(device=<default.qubit device (wires=2)",
-                ],
-            ),
             (
                 "pennylane.workflow.execution",
                 [
@@ -150,7 +144,7 @@ class TestLogging:
             ),
         ]
 
-        for expected, actual in zip(log_records_expected, caplog.records[:2]):
+        for expected, actual in zip(log_records_expected, caplog.records[:1]):
             assert expected[0] in actual.name
             for exp_msg in expected[1]:
                 assert exp_msg in actual.getMessage()
diff --git a/tests/test_qnode.py b/tests/test_qnode.py
index 68e91cd40..e3c2e8f18 100644
--- a/tests/test_qnode.py
+++ b/tests/test_qnode.py
@@ -1643,29 +1643,11 @@ class TestTransformProgramIntegration:
         assert circuit() == 100
 
 
-class TestNewDeviceIntegration:
-    """Basic tests for integration of the new device interface and the QNode."""
+class TestGetGradientFn:
+    """Test the get_gradient_fn static method."""
 
     dev = CustomDevice()
 
-    def test_initialization(self):
-        """Test that a qnode can be initialized with the new device without error."""
-
-        def f():
-            return qml.expval(qml.PauliZ(0))
-
-        qn = QNode(f, self.dev)
-        assert qn.device is self.dev
-
-    def test_repr(self):
-        """Test that the repr works with the new device."""
-
-        def f():
-            return qml.expval(qml.PauliZ(0))
-
-        qn = QNode(f, self.dev)
-        assert repr(qn) == "<QNode: device='CustomDevice', interface='auto', diff_method='best'>"
-
     def test_get_gradient_fn_custom_device(self):
         """Test get_gradient_fn is parameter for best for null device."""
         gradient_fn, kwargs, new_dev = QNode.get_gradient_fn(self.dev, "autograd", "best")
@@ -1742,6 +1724,81 @@ class TestNewDeviceIntegration:
         assert not kwargs
         assert new_dev is dev
 
+    def test_diff_method_is_none(self):
+        """Test get_gradient_fn behaves correctly."""
+        gradient_fn, kwargs, new_dev = QNode.get_gradient_fn(
+            self.dev, interface=None, diff_method=None
+        )
+        assert gradient_fn is None
+        assert not kwargs
+        assert new_dev is self.dev
+
+    def test_transform_dispatcher_as_diff_method(self):
+        """Test when diff_method is of type TransformDispatcher"""
+        gradient_fn, kwargs, new_dev = QNode.get_gradient_fn(
+            self.dev, interface=None, diff_method=qml.gradients.param_shift
+        )
+        assert gradient_fn is qml.gradients.param_shift
+        assert not kwargs
+        assert new_dev is self.dev
+
+    def test_invalid_diff_method(self):
+        """Test that an invalid diff method raises an error."""
+        with pytest.raises(
+            qml.QuantumFunctionError, match="Differentiation method invalid-method not recognized"
+        ):
+            QNode.get_gradient_fn(self.dev, None, diff_method="invalid-method")
+
+    @pytest.mark.parametrize("diff_method", ["parameter-shift", "finite-diff", "spsa", "hadamard"])
+    def test_valid_diff_method_str(self, diff_method):
+        """Test that gradient_fn are retrieved correctly."""
+        gradient_transform_map = {
+            "parameter-shift": qml.gradients.param_shift,
+            "finite-diff": qml.gradients.finite_diff,
+            "spsa": qml.gradients.spsa_grad,
+            "hadamard": qml.gradients.hadamard_grad,
+        }
+        gradient_fn, kwargs, new_dev = QNode.get_gradient_fn(
+            self.dev, interface=None, diff_method=diff_method
+        )
+        assert gradient_fn is gradient_transform_map[diff_method]
+        assert not kwargs
+        assert new_dev is self.dev
+
+    def test_param_shift_method_with_cv_ops(self):
+        """Test that 'parameter-shift-cv' is used when CV operations are present."""
+        tape = qml.tape.QuantumScript([qml.Displacement(0.5, 0.0, wires=0)])
+        gradient_fn, kwargs, new_dev = QNode.get_gradient_fn(
+            self.dev, interface=None, diff_method="parameter-shift", tape=tape
+        )
+        assert gradient_fn is qml.gradients.param_shift_cv
+        assert kwargs == {"dev": self.dev}
+        assert new_dev is self.dev
+
+
+class TestNewDeviceIntegration:
+    """Basic tests for integration of the new device interface and the QNode."""
+
+    dev = CustomDevice()
+
+    def test_initialization(self):
+        """Test that a qnode can be initialized with the new device without error."""
+
+        def f():
+            return qml.expval(qml.PauliZ(0))
+
+        qn = QNode(f, self.dev)
+        assert qn.device is self.dev
+
+    def test_repr(self):
+        """Test that the repr works with the new device."""
+
+        def f():
+            return qml.expval(qml.PauliZ(0))
+
+        qn = QNode(f, self.dev)
+        assert repr(qn) == "<QNode: device='CustomDevice', interface='auto', diff_method='best'>"
+
     def test_device_with_custom_diff_method_name(self):
         """Test a device that has its own custom diff method."""
 
diff --git a/tests/workflow/test_get_best_diff_method.py b/tests/workflow/test_get_best_diff_method.py
index 1fbe47ec9..5296a0a94 100644
--- a/tests/workflow/test_get_best_diff_method.py
+++ b/tests/workflow/test_get_best_diff_method.py
@@ -55,7 +55,7 @@ class CustomDeviceWithDiffMethod(qml.devices.Device):
 class TestValidation:
     """Tests for QNode creation and validation"""
 
-    @pytest.mark.autograd
+    @pytest.mark.jax
     def test_best_method_is_device(self):
         """Test that the method for determining the best diff method
         for a device that is a child of qml.devices.Device and has a
@@ -63,7 +63,7 @@ class TestValidation:
 
         dev = CustomDeviceWithDiffMethod()
         qn_jax = qml.QNode(dummyfunc, dev, "jax")
-        qn_none = qml.QNode(dummyfunc, dev, None)
+        qn_none = qml.QNode(dummyfunc, dev)
 
         res = get_best_diff_method(qn_jax)()
         assert res == "device"
@@ -71,6 +71,7 @@ class TestValidation:
         res = get_best_diff_method(qn_none)()
         assert res == "device"
 
+    @pytest.mark.all_interfaces
     @pytest.mark.parametrize("interface", ["jax", "tensorflow", "torch", "autograd"])
     def test_best_method_is_backprop(self, interface):
         """Test that the method for determining the best diff method