[BUG] Metric tensor fails under certain combinations of device and differentiation method

Expected behavior

I expect the code below to work with both lightning.qubit and default.qubit for any differentiation method. However qml.metric_tensor is failing for the cases outlined in the table below.

Actual behavior

The table below shows Yes if the combination of diff_method and device works or No if it doesn't. The exact error that shows up depends on the combination of device and method.

diff_method	default.qubit	lightning.qubit
backprop	Yes	No
adjoint	No	No
parameter-shift	Yes	Yes
finite-diff	Yes	Yes
hadamard	Yes	Yes
device	No	No
spsa	Yes	Yes
best	Yes	Yes

Additional information

This was originally reported as part of this forum post.

This issue might be the cause for QNGO to fail in many of these cases.

Source code

import pennylane as qml
from pennylane import numpy as pnp

# Data
data = pnp.array([0.,1.],requires_grad=False)

# Device
n_qubits=2
# We create a device with one extra wire because we need an auxiliary wire when using QNGO
dev = qml.device('default.qubit', wires=n_qubits+1)

# QNode
diff_method='adjoint'

@qml.qnode(dev,diff_method=diff_method)
def circuit(params):
  # Data embedding
  qml.RX(data[0],wires=0)
  qml.RX(data[1],wires=1)

  # Parametrized layer
  qml.Rot(params[0],params[1],params[2],wires=0)
  qml.Rot(params[0],params[1],params[2],wires=1)
  qml.Hadamard(wires=0)
  qml.CNOT(wires=[0,1])

  # Measurement
  return qml.expval(qml.Z(0))

# Initial value of the parameters
params = pnp.array([1.,2.,3.],requires_grad=True)

# Initial value of the circuit
print(circuit(params))

# Cost function
def cost_f(params):
  return pnp.abs(circuit(params))

# Optimizer
opt = qml.QNGOptimizer()

# If we're using QNGO we need to define a metric tensor function
mt_fn = qml.metric_tensor(circuit)
print(mt_fn(params))

Tracebacks

---------------------------------------------------------------------------
QuantumFunctionError                      Traceback (most recent call last)
<ipython-input-23-e113c379defd> in <cell line: 12>()
     11 
     12 @qml.qnode(dev,diff_method=diff_method)
---> 13 def circuit(params):
     14   # Data embedding
     15   qml.RX(data[0],wires=0)

2 frames
/usr/local/lib/python3.10/dist-packages/pennylane/workflow/qnode.py in __init__(self, func, device, interface, diff_method, expansion_strategy, max_expansion, grad_on_execution, cache, cachesize, max_diff, device_vjp, **gradient_kwargs)
    529 
    530         self._transform_program = qml.transforms.core.TransformProgram()
--> 531         self._update_gradient_fn()
    532         functools.update_wrapper(self, func)
    533 

/usr/local/lib/python3.10/dist-packages/pennylane/workflow/qnode.py in _update_gradient_fn(self, shots, tape)
    602             diff_method = "parameter-shift"
    603 
--> 604         self.gradient_fn, self.gradient_kwargs, self.device = self.get_gradient_fn(
    605             self._original_device, self.interface, diff_method, tape=tape
    606         )

/usr/local/lib/python3.10/dist-packages/pennylane/workflow/qnode.py in get_gradient_fn(device, interface, diff_method, tape)
    652                 return new_config.gradient_method, {}, device
    653             if diff_method in {"backprop", "adjoint", "device"}:  # device-only derivatives
--> 654                 raise qml.QuantumFunctionError(
    655                     f"Device {device} does not support {diff_method} with requested circuit."
    656                 )

QuantumFunctionError: Device <lightning.qubit device (wires=3) at 0x7b31563b2ad0> does not support backprop with requested circuit.

System information

Name: PennyLane
Version: 0.36.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, pennylane-lightning, requests, rustworkx, scipy, semantic-version, 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.25.2
Scipy version:           1.11.4
Installed devices:
- default.clifford (PennyLane-0.36.0)
- default.gaussian (PennyLane-0.36.0)
- default.mixed (PennyLane-0.36.0)
- default.qubit (PennyLane-0.36.0)
- default.qubit.autograd (PennyLane-0.36.0)
- default.qubit.jax (PennyLane-0.36.0)
- default.qubit.legacy (PennyLane-0.36.0)
- default.qubit.tf (PennyLane-0.36.0)
- default.qubit.torch (PennyLane-0.36.0)
- default.qutrit (PennyLane-0.36.0)
- default.qutrit.mixed (PennyLane-0.36.0)
- null.qubit (PennyLane-0.36.0)
- lightning.qubit (PennyLane_Lightning-0.36.0)

Existing GitHub issues

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

PennyLaneAI / pennylane

[BUG] Metric tensor fails under certain combinations of device and differentiation method #5901