Support for `qml.state` when `diff_method="adjoint"`

isaacdevlugt commented 1 month ago

Issue description

Lightning does not support measuring qml.state when diff_method="adjoint".

Expected behavior: QNodes running on lightning can return qml.state when diff_method="adjoint" (also for circuits with qml.Snapshot())
Actual behavior: Raises an error.
Reproduces how often: 100%
System information:

Name: PennyLane
Version: 0.37.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: [/Users/isaac/.virtualenvs/pennylane-catalyst/lib/python3.11/site-packages](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/.virtualenvs/pennylane-catalyst/lib/python3.11/site-packages)
Requires: appdirs, autograd, autoray, cachetools, networkx, numpy, packaging, pennylane-lightning, requests, rustworkx, scipy, semantic-version, toml, typing-extensions
Required-by: PennyLane-Catalyst, PennyLane_Lightning

Platform info:           macOS-14.5-arm64-arm-64bit
Python version:          3.11.8
Numpy version:           1.26.4
Scipy version:           1.12.0
Installed devices:
- default.clifford (PennyLane-0.38.0.dev0)
- default.gaussian (PennyLane-0.38.0.dev0)
- default.mixed (PennyLane-0.38.0.dev0)
- default.qubit (PennyLane-0.38.0.dev0)
- default.qubit.autograd (PennyLane-0.38.0.dev0)
- default.qubit.jax (PennyLane-0.38.0.dev0)
- default.qubit.legacy (PennyLane-0.38.0.dev0)
- default.qubit.tf (PennyLane-0.38.0.dev0)
- default.qubit.torch (PennyLane-0.38.0.dev0)
- default.qutrit (PennyLane-0.38.0.dev0)
- default.qutrit.mixed (PennyLane-0.38.0.dev0)
- default.tensor (PennyLane-0.38.0.dev0)
- null.qubit (PennyLane-0.38.0.dev0)
- lightning.qubit (PennyLane_Lightning-0.37.0)
- nvidia.custatevec (PennyLane-Catalyst-0.7.0)
- nvidia.cutensornet (PennyLane-Catalyst-0.7.0)
- oqc.cloud (PennyLane-Catalyst-0.7.0)
- softwareq.qpp (PennyLane-Catalyst-0.7.0)

Source code and tracebacks

QNode return:

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

@qml.qnode(dev, diff_method='adjoint')
def circuit():
    qml.Hadamard(0)
    return qml.state()

circuit()

---------------------------------------------------------------------------
QuantumFunctionError                      Traceback (most recent call last)
Cell In[12], [line 10](vscode-notebook-cell:?execution_count=12&line=10)
      [6](vscode-notebook-cell:?execution_count=12&line=6)     #qml.Snapshot()
      [7](vscode-notebook-cell:?execution_count=12&line=7)     #return qml.expval(qml.Z(0))
      [8](vscode-notebook-cell:?execution_count=12&line=8)     return qml.state()
---> [10](vscode-notebook-cell:?execution_count=12&line=10) circuit()
     [11](vscode-notebook-cell:?execution_count=12&line=11) #qml.snapshots(circuit)()

File ~/Documents/pennylane/pennylane/workflow/qnode.py:1164, in QNode.__call__(self, *args, **kwargs)
   [1162](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/workflow/qnode.py:1162) if qml.capture.enabled():
   [1163](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/workflow/qnode.py:1163)     return qml.capture.qnode_call(self, *args, **kwargs)
-> [1164](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/workflow/qnode.py:1164) return self._impl_call(*args, **kwargs)

File ~/Documents/pennylane/pennylane/workflow/qnode.py:1147, in QNode._impl_call(self, *args, **kwargs)
   [1144](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/workflow/qnode.py:1144) self.construct(args, kwargs)
   [1146](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/workflow/qnode.py:1146) original_grad_fn = [self.gradient_fn, self.gradient_kwargs, self.device]
-> [1147](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/workflow/qnode.py:1147) self._update_gradient_fn(shots=override_shots, tape=self._tape)
   [1149](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/workflow/qnode.py:1149) try:
   [1150](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/workflow/qnode.py:1150)     res = self._execution_component(args, kwargs, override_shots=override_shots)

File ~/Documents/pennylane/pennylane/workflow/qnode.py:632, in QNode._update_gradient_fn(self, shots, tape)
    [625](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/workflow/qnode.py:625) if (
    [626](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/workflow/qnode.py:626)     self.device.name == "lightning.qubit"
    [627](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/workflow/qnode.py:627)     and qml.metric_tensor in self.transform_program
    [628](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/workflow/qnode.py:628)     and self.diff_method == "best"
    [629](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/workflow/qnode.py:629) ):
    [630](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/workflow/qnode.py:630)     diff_method = "parameter-shift"
--> [632](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/workflow/qnode.py:632) self.gradient_fn, self.gradient_kwargs, self.device = QNode.get_gradient_fn(
    [633](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/workflow/qnode.py:633)     self._original_device, self.interface, diff_method, tape=tape
    [634](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/workflow/qnode.py:634) )
    [635](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/workflow/qnode.py:635) self.gradient_kwargs.update(self._user_gradient_kwargs or {})

File ~/Documents/pennylane/pennylane/logging/decorators.py:61, in log_string_debug_func.<locals>.wrapper_entry(*args, **kwargs)
     [54](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/logging/decorators.py:54)     s_caller = "::L".join(
     [55](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/logging/decorators.py:55)         [str(i) for i in inspect.getouterframes(inspect.currentframe(), 2)[1][1:3]]
     [56](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/logging/decorators.py:56)     )
     [57](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/logging/decorators.py:57)     lgr.debug(
     [58](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/logging/decorators.py:58)         f"Calling {f_string} from {s_caller}",
     [59](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/logging/decorators.py:59)         **_debug_log_kwargs,
     [60](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/logging/decorators.py:60)     )
---> [61](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/logging/decorators.py:61) return func(*args, **kwargs)

File ~/Documents/pennylane/pennylane/workflow/qnode.py:683, in QNode.get_gradient_fn(device, interface, diff_method, tape)
    [681](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/workflow/qnode.py:681)         return new_config.gradient_method, {}, device
    [682](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/workflow/qnode.py:682)     if diff_method in {"backprop", "adjoint", "device"}:  # device-only derivatives
--> [683](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/workflow/qnode.py:683)         raise qml.QuantumFunctionError(
    [684](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/workflow/qnode.py:684)             f"Device {device} does not support {diff_method} with requested circuit."
    [685](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/workflow/qnode.py:685)         )
    [687](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/workflow/qnode.py:687) if diff_method == "best":
    [688](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/workflow/qnode.py:688)     return QNode.get_best_method(device, interface, tape=tape)

QuantumFunctionError: Device <lightning.qubit device (wires=2) at 0x1224cef50> does not support adjoint with requested circuit.

Snapshots

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

@qml.qnode(dev, diff_method='adjoint')
def circuit():
    qml.Hadamard(0)
    qml.Snapshot()
    return qml.expval(qml.Z(0))

qml.snapshots(circuit)()

[/Users/isaac/Documents/pennylane/pennylane/debugging/snapshot.py:247](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/pennylane/debugging/snapshot.py:247): UserWarning: Snapshots are not supported for the given device. Therefore, a tape will be created for each snapshot, resulting in a total of n_snapshots + 1 executions.
  warnings.warn(
---------------------------------------------------------------------------
DeviceError                               Traceback (most recent call last)
Cell In[15], [line 10](vscode-notebook-cell:?execution_count=15&line=10)
      [7](vscode-notebook-cell:?execution_count=15&line=7)     return qml.expval(qml.Z(0))
      [9](vscode-notebook-cell:?execution_count=15&line=9) #circuit()
---> [10](vscode-notebook-cell:?execution_count=15&line=10) qml.snapshots(circuit)()

File ~/Documents/pennylane/pennylane/workflow/qnode.py:1164, in QNode.__call__(self, *args, **kwargs)
   [1162](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/workflow/qnode.py:1162) if qml.capture.enabled():
   [1163](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/workflow/qnode.py:1163)     return qml.capture.qnode_call(self, *args, **kwargs)
-> [1164](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/workflow/qnode.py:1164) return self._impl_call(*args, **kwargs)

File ~/Documents/pennylane/pennylane/workflow/qnode.py:1150, in QNode._impl_call(self, *args, **kwargs)
   [1147](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/workflow/qnode.py:1147) self._update_gradient_fn(shots=override_shots, tape=self._tape)
   [1149](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/workflow/qnode.py:1149) try:
-> [1150](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/workflow/qnode.py:1150)     res = self._execution_component(args, kwargs, override_shots=override_shots)
   [1151](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/workflow/qnode.py:1151) finally:
   [1152](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/workflow/qnode.py:1152)     if old_interface == "auto":

File ~/Documents/pennylane/pennylane/workflow/qnode.py:1103, in QNode._execution_component(self, args, kwargs, override_shots)
   [1100](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/workflow/qnode.py:1100) _prune_dynamic_transform(full_transform_program, inner_transform_program)
   [1102](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/workflow/qnode.py:1102) # pylint: disable=unexpected-keyword-arg
-> [1103](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/workflow/qnode.py:1103) res = qml.execute(
   [1104](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/workflow/qnode.py:1104)     (self._tape,),
   [1105](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/workflow/qnode.py:1105)     device=self.device,
   [1106](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/workflow/qnode.py:1106)     gradient_fn=self.gradient_fn,
   [1107](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/workflow/qnode.py:1107)     interface=self.interface,
   [1108](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/workflow/qnode.py:1108)     transform_program=full_transform_program,
   [1109](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/workflow/qnode.py:1109)     inner_transform=inner_transform_program,
   [1110](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/workflow/qnode.py:1110)     config=config,
   [1111](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/workflow/qnode.py:1111)     gradient_kwargs=self.gradient_kwargs,
   [1112](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/workflow/qnode.py:1112)     override_shots=override_shots,
   [1113](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/workflow/qnode.py:1113)     **self.execute_kwargs,
   [1114](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/workflow/qnode.py:1114) )
   [1115](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/workflow/qnode.py:1115) res = res[0]
   [1117](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/workflow/qnode.py:1117) # convert result to the interface in case the qfunc has no parameters

File ~/Documents/pennylane/pennylane/workflow/execution.py:650, in execute(tapes, device, gradient_fn, interface, transform_program, inner_transform, config, grad_on_execution, gradient_kwargs, cache, cachesize, max_diff, override_shots, expand_fn, max_expansion, device_batch_transform, device_vjp, mcm_config)
    [645](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/workflow/execution.py:645)     if not device_batch_transform:
    [646](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/workflow/execution.py:646)         warnings.warn(
    [647](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/workflow/execution.py:647)             "device batch transforms cannot be turned off with the new device interface.",
    [648](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/workflow/execution.py:648)             UserWarning,
    [649](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/workflow/execution.py:649)         )
--> [650](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/workflow/execution.py:650)     tapes, post_processing = transform_program(tapes)
    [651](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/workflow/execution.py:651) else:
    [652](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/workflow/execution.py:652)     # TODO: Remove once old device are removed
    [653](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/workflow/execution.py:653)     tapes, program_post_processing = transform_program(tapes)

File ~/Documents/pennylane/pennylane/transforms/core/transform_program.py:515, in TransformProgram.__call__(self, tapes)
    [513](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/transforms/core/transform_program.py:513) if self._argnums is not None and self._argnums[i] is not None:
    [514](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/transforms/core/transform_program.py:514)     tape.trainable_params = self._argnums[i][j]
--> [515](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/transforms/core/transform_program.py:515) new_tapes, fn = transform(tape, *targs, **tkwargs)
    [516](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/transforms/core/transform_program.py:516) execution_tapes.extend(new_tapes)
    [518](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/transforms/core/transform_program.py:518) fns.append(fn)

File ~/Documents/pennylane/pennylane/devices/preprocess.py:492, in validate_measurements(tape, analytic_measurements, sample_measurements, name)
    [490](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/devices/preprocess.py:490)     for m in chain(snapshot_measurements, tape.measurements):
    [491](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/devices/preprocess.py:491)         if not analytic_measurements(m):
--> [492](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/devices/preprocess.py:492)             raise DeviceError(
    [493](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/devices/preprocess.py:493)                 f"Measurement {m} not accepted for analytic simulation on {name}."
    [494](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/devices/preprocess.py:494)             )
    [496](https://file+.vscode-resource.vscode-cdn.net/Users/isaac/Documents/pennylane/~/Documents/pennylane/pennylane/devices/preprocess.py:496) return (tape,), null_postprocessing

DeviceError: Measurement state(wires=[0, 1]) not accepted for analytic simulation on adjoint + lightning.qubit.

sotskopa commented 1 month ago

I have the same problem with circuits returning qml.probs.

isaacdevlugt commented 1 month ago

I have the same problem with circuits returning qml.probs.

Hey @sotskopa, thank you for letting us know! FYI, you can bypass the issue by setting a different differentiation method (e.g. diff_method="parameter-shift")

isaacdevlugt commented 1 month ago

Just leaving a note here after some internal discussions about this issue. We won't be fixing this bug at this time and will leave it as a future endeavour.

Re: the error message for state-based MPs + Snapshot + adjoint, it's not an easy fix to have a message that better indicates what's going on (e.g., lightning + adjoint doesn't support state-based measurements with Snapshots instead of the error message above) because of the order in which things are executed under the hood.

trbromley commented 1 month ago

FYI, you can bypass the issue by setting a different differentiation method (e.g. diff_method="parameter-shift")

Note that lightning.qubit doesn't currently support differentiating the state with any differentiation method, but you can use diff_method="adjoint" or diff_method="backprop" with default.qubit:

import pennylane as qml
from jax import numpy as np
import jax

dev = qml.device("default.qubit", wires=1)

@qml.qnode(dev, diff_method="adjoint")
def f(x):
    qml.RX(x, 0)
    return qml.state()

x = np.array(0.4, dtype=np.complex64)

jax.jacobian(f, holomorphic=True)(x)

PennyLaneAI / pennylane-lightning

Support for `qml.state` when `diff_method="adjoint"` #808

Issue description

Source code and tracebacks