Update logic for enabling grad_on_execution in lightning_*.py devices (#1016)

**Context:**

In PennyLane, `grad_on_execution` is expected to be enabled when the
gradient method is set to `adjoint`. Since this is automatically true
for Lightning devices, the `use_device_gradient` flag can serve as an
additional condition to decide whether to enable `grad_on_execution`.

This discrepancy caused certain Lightning device tests to fail during a
PennyLane workflow clean-up PR. The issue arose because PennyLane
expected `grad_on_execution` to be `None` but received `True` instead.
Further details can be found in the [PR
workflow](https://github.com/PennyLaneAI/pennylane/actions/runs/12129490447/job/33818035152?pr=6657).

**Description of the Change:**

Update logic for when `grad_on_execution` is turned on in `lightning`
device configurations.

**Benefits:** Unblocks
PennyLaneAI/pennylane#6657

**Possible Drawbacks:** None.

[sc-72150]

---------

Co-authored-by: ringo-but-quantum <[email protected]>
Co-authored-by: Ali Asadi <[email protected]>
Co-authored-by: Christina Lee <[email protected]>

Author: 4 people

.github/CHANGELOG.md

@@ -26,6 +26,9 @@
 
 ### Improvements
 
+* Update the logic for enabling `grad_on_execution` during device execution.
+  [(#1016)](https://github.com/PennyLaneAI/pennylane-lightning/pull/1016)
+
 * Reverse Lightning Qubit generators vector insertion order.
   [(#1009)](https://github.com/PennyLaneAI/pennylane-lightning/pull/1009)
 

pennylane_lightning/core/_version.py

@@ -16,4 +16,4 @@
    Version number (major.minor.patch[-label])
 """
 
-__version__ = "0.40.0-dev26"
+__version__ = "0.40.0-dev27"

pennylane_lightning/lightning_gpu/lightning_gpu.py

@@ -299,7 +299,11 @@ def _setup_execution_config(self, config):
             updated_values["gradient_method"] = "adjoint"
         if config.use_device_gradient is None:
             updated_values["use_device_gradient"] = config.gradient_method in ("best", "adjoint")
-        if config.grad_on_execution is None:
+        if (
+            config.use_device_gradient
+            or updated_values.get("use_device_gradient", False)
+            and config.grad_on_execution is None
+        ):
             updated_values["grad_on_execution"] = True
 
         new_device_options = dict(config.device_options)

pennylane_lightning/lightning_kokkos/lightning_kokkos.py

@@ -271,7 +271,11 @@ def _setup_execution_config(self, config):
             updated_values["gradient_method"] = "adjoint"
         if config.use_device_gradient is None:
             updated_values["use_device_gradient"] = config.gradient_method in ("best", "adjoint")
-        if config.grad_on_execution is None:
+        if (
+            config.use_device_gradient
+            or updated_values.get("use_device_gradient", False)
+            and config.grad_on_execution is None
+        ):
             updated_values["grad_on_execution"] = True
 
         new_device_options = dict(config.device_options)

pennylane_lightning/lightning_qubit/lightning_qubit.py

@@ -310,7 +310,11 @@ def _setup_execution_config(self, config):
             updated_values["gradient_method"] = "adjoint"
         if config.use_device_gradient is None:
             updated_values["use_device_gradient"] = config.gradient_method in ("best", "adjoint")
-        if config.grad_on_execution is None:
+        if (
+            config.use_device_gradient
+            or updated_values.get("use_device_gradient", False)
+            and config.grad_on_execution is None
+        ):
             updated_values["grad_on_execution"] = True
 
         new_device_options = dict(config.device_options)