[BUG] QNSPSAOptimizer improperly handles optimization on lightning.qubit

[isaacdevlugt]

Expected behavior

No difference in behaviour with default.qubit and lightning.qubit when optimizing a circuit with QNSPSAOptimizer.

Actual behavior

QNSPSAOptimizer mishandles the case where lightning qubit is used.

Additional information

Might be related to #5437? Not sure.

Source code

import pennylane as qml
from pennylane import numpy as np

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

@qml.qnode(dev)
def cost(params):
    qml.RX(params[0], wires=0)
    qml.CRY(params[1], wires=[0, 1])
    return qml.expval(qml.Z(0) @ qml.Z(1))

params = np.random.rand(2)
opt = qml.QNSPSAOptimizer(stepsize=5e-2)

print(type(cost(params)))

for _ in range(2):
    print(type(params))
    params, loss = opt.step_and_cost(cost, params)

Tracebacks

<class 'numpy.ndarray'>
<class 'pennylane.numpy.tensor.tensor'>
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Cell In[10], line 19
     17 for _ in range(2):
     18     print(type(params))
---> 19     params, loss = opt.step_and_cost(cost, params)

File ~/.virtualenvs/pennylane-stable/lib/python3.11/site-packages/pennylane/optimize/qnspsa.py:191, in QNSPSAOptimizer.step_and_cost(self, cost, *args, **kwargs)
    188     loss_curr = cost(*args, **kwargs)
    189     return params_next, loss_curr
--> 191 params_next, loss_curr = self._apply_blocking(cost, args, kwargs, params_next)
    192 return params_next, loss_curr

File ~/.virtualenvs/pennylane-stable/lib/python3.11/site-packages/pennylane/optimize/qnspsa.py:441, in QNSPSAOptimizer._apply_blocking(self, cost, args, kwargs, params_next)
    439 cost.construct(params_next, kwargs)
    440 tape_loss_next = cost.tape.copy(copy_operations=True)
--> 441 program, _ = cost.device.preprocess()
    442 loss_curr, loss_next = qml.execute(
    443     [tape_loss_curr, tape_loss_next], cost.device, None, transform_program=program
    444 )
    446 # self.k has been updated earlier

AttributeError: 'LightningQubit' object has no attribute 'preprocess'

System information

Name: PennyLane
Version: 0.35.1
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-stable/lib/python3.11/site-packages
Requires: appdirs, autograd, autoray, cachetools, networkx, numpy, pennylane-lightning, requests, rustworkx, scipy, semantic-version, toml, typing-extensions
Required-by: PennyLane_Lightning

Platform info:           macOS-14.4.1-arm64-arm-64bit
Python version:          3.11.8
Numpy version:           1.26.4
Scipy version:           1.12.0
Installed devices:
- default.clifford (PennyLane-0.35.1)
- default.gaussian (PennyLane-0.35.1)
- default.mixed (PennyLane-0.35.1)
- default.qubit (PennyLane-0.35.1)
- default.qubit.autograd (PennyLane-0.35.1)
- default.qubit.jax (PennyLane-0.35.1)
- default.qubit.legacy (PennyLane-0.35.1)
- default.qubit.tf (PennyLane-0.35.1)
- default.qubit.torch (PennyLane-0.35.1)
- default.qutrit (PennyLane-0.35.1)
- null.qubit (PennyLane-0.35.1)
- lightning.qubit (PennyLane_Lightning-0.35.1)

Existing GitHub issues

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

[EuGig]

Hi @isaacdevlugt, the problem seems to not concern QNSPSAOptimizer . The same AttributeError is returned when calling every function defined in the class LightningQubit, even if no Optimizer is defined:

>>> import pennylane as qml
>>> dev = qml.device("lightning.qubit", wires=2)
>>> dev._setup_execution_config()
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
AttributeError: 'LightningQubit' object has no attribute '_setup_execution_config'
>>>
>>> dev.preprocess()
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
AttributeError: 'LightningQubit' object has no attribute 'preprocess'
>>>
>>> dev.support_derivatives()
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
AttributeError: 'LightningQubit' object has no attribute 'support_derivatives'

The same happens with functions that take a circuit as an argument.
In fact, using help(dev) does not show those methods.
Maybe the problem is not the optimizer but how LightningQubit is imported is some __init__.py file. Even the order of how things are imported matters. I can dig more into it.

[albi3ro]

The issue is that Line 441 strictly assumes the new device interface. While LightningQubit does follow the new device interface as of two weeks ago, we do still need to support the old device interface with QNSPSA. This should be tested with default.qubit.legacy as well as just default.qubit.

[PietropaoloFrisoni]

Hi everyone!

I was going to fix the bug, but when I tried to run the above code, I didn't see an error. Can you please confirm if the issue is still present?

Thanks!

[isaacdevlugt]

@PietropaoloFrisoni we discussed internally, but the issue is no longer present with lightning because lightning now uses the new device API (PennyLaneAI/pennylane-lightning#646). But, QNSPSAOptimizer needs to be tested with default.qubit.legacy

[PietropaoloFrisoni]

@isaacdevlugt Thanks, Isaac! That's precisely what I did to reproduce and fix the issue (you can see it in the description of #5497 )