[BUG] diff_method=None sometimes returns a result and sometimes an error

[CatalinaAlbornoz]

Expected behavior

The docs say that when someone uses diff_method=None the QNode cannot be differentiated.

Therefore I expect that when someone tries to differentiate a qnode that has diff_method=None they should get an error message saying something along the lines of "A QNode with diff_method=None cannot be differentiated".

At the very least I expect a note in the docs in the Gradients and Training section explaining when it's expected to work and what error arises when it doesn't work.

Actual behavior

In practice when using default.qubit we do return a result while lightning.qubit returns a rather incomprehensible error.

Additional information

It's been discussed that changing behaviour here might cause breaking changes. We should take some time to figure out whether the best solution is in the code or in the docs.

Source code

import pennylane as qml
from pennylane import numpy as pnp

dev = qml.device('lightning.qubit', wires=2)

@qml.qnode(dev, diff_method=None)
def circuit4(phi, theta):
    qml.RX(phi[0], wires=0)
    qml.RZ(phi[1], wires=1)
    qml.CNOT(wires=[0, 1])
    qml.RX(theta, wires=0)
    return qml.expval(qml.PauliZ(0))

phi = pnp.array([0.2, 1.0])
theta = 0.2

print(qml.grad(circuit4)(phi, theta))

Tracebacks

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-34-2c18c638b76d> in <cell line: 20>()
     18 theta = 0.2
     19 
---> 20 print(qml.grad(circuit4)(phi, theta))

17 frames
/usr/local/lib/python3.10/dist-packages/pennylane/_grad.py in __call__(self, *args, **kwargs)
    222             return ()
    223 
--> 224         grad_value, ans = grad_fn(*args, **kwargs)  # pylint: disable=not-callable
    225         self._forward = ans
    226 

/usr/local/lib/python3.10/dist-packages/autograd/wrap_util.py in nary_f(*args, **kwargs)
     18             else:
     19                 x = tuple(args[i] for i in argnum)
---> 20             return unary_operator(unary_f, x, *nary_op_args, **nary_op_kwargs)
     21         return nary_f
     22     return nary_operator

/usr/local/lib/python3.10/dist-packages/pennylane/_grad.py in _grad_with_forward(fun, x)
    240         difference being that it returns both the gradient *and* the forward pass
    241         value."""
--> 242         vjp, ans = _make_vjp(fun, x)  # pylint: disable=redefined-outer-name
    243 
    244         if vspace(ans).size != 1:

/usr/local/lib/python3.10/dist-packages/autograd/core.py in make_vjp(fun, x)
      8 def make_vjp(fun, x):
      9     start_node = VJPNode.new_root()
---> 10     end_value, end_node =  trace(start_node, fun, x)
     11     if end_node is None:
     12         def vjp(g): return vspace(x).zeros()

/usr/local/lib/python3.10/dist-packages/autograd/tracer.py in trace(start_node, fun, x)
      8     with trace_stack.new_trace() as t:
      9         start_box = new_box(x, t, start_node)
---> 10         end_box = fun(start_box)
     11         if isbox(end_box) and end_box._trace == start_box._trace:
     12             return end_box._value, end_box._node

/usr/local/lib/python3.10/dist-packages/autograd/wrap_util.py in unary_f(x)
     13                 else:
     14                     subargs = subvals(args, zip(argnum, x))
---> 15                 return fun(*subargs, **kwargs)
     16             if isinstance(argnum, int):
     17                 x = args[argnum]

/usr/local/lib/python3.10/dist-packages/pennylane/workflow/qnode.py in __call__(self, *args, **kwargs)
    985         if qml.capture.enabled():
    986             return qml.capture.qnode_call(self, *args, **kwargs)
--> 987         return self._impl_call(*args, **kwargs)
    988 
    989 

/usr/local/lib/python3.10/dist-packages/pennylane/workflow/qnode.py in _impl_call(self, *args, **kwargs)
    975 
    976         try:
--> 977             res = self._execution_component(args, kwargs)
    978         finally:
    979             if old_interface == "auto":

/usr/local/lib/python3.10/dist-packages/pennylane/workflow/qnode.py in _execution_component(self, args, kwargs)
    933 
    934         # pylint: disable=unexpected-keyword-arg
--> 935         res = qml.execute(
    936             (self._tape,),
    937             device=self.device,

/usr/local/lib/python3.10/dist-packages/pennylane/workflow/execution.py in execute(tapes, device, gradient_fn, interface, transform_program, inner_transform, config, grad_on_execution, gradient_kwargs, cache, cachesize, max_diff, device_vjp, mcm_config)
    521     # Exiting early if we do not need to deal with an interface boundary
    522     if no_interface_boundary_required:
--> 523         results = inner_execute(tapes)
    524         return post_processing(results)
    525 

/usr/local/lib/python3.10/dist-packages/pennylane/workflow/execution.py in inner_execute(tapes, **_)
    200 
    201         if transformed_tapes:
--> 202             results = device.execute(transformed_tapes, execution_config=execution_config)
    203         else:
    204             results = ()

/usr/local/lib/python3.10/dist-packages/pennylane/devices/modifiers/simulator_tracking.py in execute(self, circuits, execution_config)
     28     @wraps(untracked_execute)
     29     def execute(self, circuits, execution_config=DefaultExecutionConfig):
---> 30         results = untracked_execute(self, circuits, execution_config)
     31         if isinstance(circuits, QuantumScript):
     32             batch = (circuits,)

/usr/local/lib/python3.10/dist-packages/pennylane/devices/modifiers/single_tape_support.py in execute(self, circuits, execution_config)
     30             is_single_circuit = True
     31             circuits = (circuits,)
---> 32         results = batch_execute(self, circuits, execution_config)
     33         return results[0] if is_single_circuit else results
     34 

/usr/local/lib/python3.10/dist-packages/pennylane_lightning/lightning_qubit/lightning_qubit.py in execute(self, circuits, execution_config)
    476                 [circuit], _ = qml.map_wires(circuit, self._wire_map)
    477             results.append(
--> 478                 self.simulate(
    479                     circuit,
    480                     self._statevector,

/usr/local/lib/python3.10/dist-packages/pennylane_lightning/lightning_qubit/lightning_qubit.py in simulate(self, circuit, state, mcmc, postselect_mode)
    559 
    560         state.reset_state()
--> 561         final_state = state.get_final_state(circuit)
    562         return self.LightningMeasurements(final_state, **mcmc).measure_final_state(circuit)

/usr/local/lib/python3.10/dist-packages/pennylane_lightning/core/_state_vector_base.py in get_final_state(self, circuit, mid_measurements, postselect_mode)
    214 
    215         """
--> 216         self.apply_operations(
    217             circuit.operations, mid_measurements=mid_measurements, postselect_mode=postselect_mode
    218         )

/usr/local/lib/python3.10/dist-packages/pennylane_lightning/core/_state_vector_base.py in apply_operations(self, operations, mid_measurements, postselect_mode)
    188                 self._apply_basis_state(operations[0].parameters[0], operations[0].wires)
    189                 operations = operations[1:]
--> 190         self._apply_lightning(
    191             operations, mid_measurements=mid_measurements, postselect_mode=postselect_mode
    192         )

/usr/local/lib/python3.10/dist-packages/pennylane_lightning/lightning_qubit/_state_vector.py in _apply_lightning(self, operations, mid_measurements, postselect_mode)
    213             elif method is not None:  # apply specialized gate
    214                 param = operation.parameters
--> 215                 method(wires, invert_param, param)
    216             elif isinstance(operation, qml.ops.Controlled):  # apply n-controlled gate
    217                 self._apply_lightning_controlled(operation)

TypeError: RX(): incompatible function arguments. The following argument types are supported:
    1. (self: pennylane_lightning.lightning_qubit_ops.StateVectorC128, arg0: list[int], arg1: bool, arg2: list[float]) -> None
    2. (self: pennylane_lightning.lightning_qubit_ops.StateVectorC128, arg0: list[int], arg1: list[bool], arg2: list[int], arg3: bool, arg4: list[float]) -> None

Invoked with: <pennylane_lightning.lightning_qubit_ops.StateVectorC128 object at 0x7a52dcdc15b0>, [0], False, [<autograd.numpy.numpy_boxes.ArrayBox object at 0x7a52dd438880>]

System information

Name: PennyLane
Version: 0.39.0
Summary: PennyLane is a cross-platform Python library for quantum computing, quantum machine learning, and quantum chemistry. Train a quantum computer the same way as a neural network.
Home-page: https://github.com/PennyLaneAI/pennylane
Author: 
Author-email: 
License: Apache License 2.0
Location: /usr/local/lib/python3.10/dist-packages
Requires: appdirs, autograd, autoray, cachetools, networkx, numpy, packaging, pennylane-lightning, requests, rustworkx, scipy, toml, typing-extensions
Required-by: PennyLane_Lightning

Platform info:           Linux-6.1.85+-x86_64-with-glibc2.35
Python version:          3.10.12
Numpy version:           1.26.4
Scipy version:           1.13.1
Installed devices:
- default.clifford (PennyLane-0.39.0)
- default.gaussian (PennyLane-0.39.0)
- default.mixed (PennyLane-0.39.0)
- default.qubit (PennyLane-0.39.0)
- default.qutrit (PennyLane-0.39.0)
- default.qutrit.mixed (PennyLane-0.39.0)
- default.tensor (PennyLane-0.39.0)
- null.qubit (PennyLane-0.39.0)
- reference.qubit (PennyLane-0.39.0)
- lightning.qubit (PennyLane_Lightning-0.39.0)

Existing GitHub issues

• I have searched existing GitHub issues to make sure the issue does not already exist.