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

[CatalinaAlbornoz]

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

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

[albi3ro]

The metric tensor converts to measuring probabilities, and adjoint is only compatible with measuring expectation values.

lightning does not support backprop at all, as expected.

As for the issue posted above, this is the same issue we encountered last release cycle in #5624. We added a lightning specific patch, but obviously this issue is still an issue and now causing problems with default.qubit too.

[trbromley]

adjoint is only compatible with measuring expectation values

@albi3ro does default.qubit support probability-based returns with the adjoint diff method?

[albi3ro]

So the issue with default.qubit is that the metric_tensor is a final transform, and final transforms run after device preprocessing. And so since the metric tensor happens after device preprocessing, it converts the batch into one no longer supported by the device.

Potentially @rmoyard might shed some light on why we can't call metric_tensor before device preprocessing?

[rmoyard]

It is a final transform because no other user transform is compatible with it but it might work with the device preprocessing. It is worth giving it a try.

[albi3ro]

What would be the difference between a user decomposing to a target gateset and a device decomposing to the target gateset? Or splitting up the measurements? They do the same thing regardless of who is requesting the transform is applied. Though I guess I'm still a little confused about why it isn't compatible with other transforms. Does it have to do with the classical jacobian component?