[BUG] Exception when deepcopy because of non-pickle-able lightning device

TheRisenPhoenix commented 5 months ago

Issue description

I use a pennylane circuit inside a torch model, which needs to be (deep)copied. If I use the default.qubit device, everything works fine, but when I switch to a lightning device (lightning.qubit / lightning.gpu), an exception is thrown. I provide a minimal example to reproduce below.

Expected behavior: (What you expect to happen) I hope to be able to copy (pickle) models regardless of the device they're running on. If not, then default.qubit and lightning.qubit should at least show the same behavior (or maybe fail with an explanation, if it is indeed intended to not clone models)
Actual behavior: (What actually happens) I get an error message, stating "TypeError: cannot pickle 'pennylane_lightning.lightning_gpu_ops.DevPool' object" (or respectively the same for pennylane_lightning.lightning_qubit_ops.StateVectorC128 )
Reproduces how often: (What percentage of the time does it reproduce?) Every time, on different machines

System information: (post the output of import pennylane as qml; qml.about())


>>> import pennylane as qml; qml.about()
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: ###mypath###/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, PennyLane_Lightning_GPU

Platform info: Linux-5.15.146.1-microsoft-standard-WSL2-x86_64-with-glibc2.35 Python version: 3.11.7 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.gpu (PennyLane_Lightning_GPU-0.35.1)
lightning.qubit (PennyLane_Lightning-0.35.1)

Source code and tracebacks

I've prepared this example which I hope helps to reproduce the issue:

import pennylane as qml
import torch
from copy import deepcopy

n_qubits = 2
n_layers = 6
weight_shapes = {"weights": (n_layers, n_qubits)}

# default.qubit works fine
# lightning.qubit also fails
dev = qml.device("lightning.gpu", wires=n_qubits)

@qml.qnode(dev)
def qnode(inputs, weights):
    qml.AngleEmbedding(inputs, wires=range(n_qubits))
    qml.BasicEntanglerLayers(weights, wires=range(n_qubits))
    return [qml.expval(qml.PauliZ(wires=i)) for i in range(n_qubits)]

class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.qlayer_1 = qml.qnn.TorchLayer(qnode, weight_shapes)

    def forward(self, x):
        x = self.qlayer_1(x)
        return x

model = Model()
model2 = deepcopy(model)

This is the stacktrace:

Traceback (most recent call last):
  File "/mnt/d/Code/report.py", line 32, in <module>
    model2 = deepcopy(model)
             ^^^^^^^^^^^^^^^
  File "/lib/python3.11/copy.py", line 172, in deepcopy
    y = _reconstruct(x, memo, *rv)
        ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/lib/python3.11/copy.py", line 271, in _reconstruct
    state = deepcopy(state, memo)
            ^^^^^^^^^^^^^^^^^^^^^
  File "/lib/python3.11/copy.py", line 146, in deepcopy
    y = copier(x, memo)
        ^^^^^^^^^^^^^^^
  File "/lib/python3.11/copy.py", line 231, in _deepcopy_dict
    y[deepcopy(key, memo)] = deepcopy(value, memo)
                             ^^^^^^^^^^^^^^^^^^^^^
  File "/lib/python3.11/copy.py", line 172, in deepcopy
    y = _reconstruct(x, memo, *rv)
        ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/lib/python3.11/copy.py", line 297, in _reconstruct
    value = deepcopy(value, memo)
            ^^^^^^^^^^^^^^^^^^^^^
  File "/lib/python3.11/copy.py", line 172, in deepcopy
    y = _reconstruct(x, memo, *rv)
        ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/lib/python3.11/copy.py", line 271, in _reconstruct
    state = deepcopy(state, memo)
            ^^^^^^^^^^^^^^^^^^^^^
  File "/lib/python3.11/copy.py", line 146, in deepcopy
    y = copier(x, memo)
        ^^^^^^^^^^^^^^^
  File "/lib/python3.11/copy.py", line 231, in _deepcopy_dict
    y[deepcopy(key, memo)] = deepcopy(value, memo)
                             ^^^^^^^^^^^^^^^^^^^^^
  File "/lib/python3.11/copy.py", line 172, in deepcopy
    y = _reconstruct(x, memo, *rv)
        ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/lib/python3.11/copy.py", line 271, in _reconstruct
    state = deepcopy(state, memo)
            ^^^^^^^^^^^^^^^^^^^^^
  File "/lib/python3.11/copy.py", line 146, in deepcopy
    y = copier(x, memo)
        ^^^^^^^^^^^^^^^
  File "/lib/python3.11/copy.py", line 231, in _deepcopy_dict
    y[deepcopy(key, memo)] = deepcopy(value, memo)
                             ^^^^^^^^^^^^^^^^^^^^^
  File "/lib/python3.11/copy.py", line 172, in deepcopy
    y = _reconstruct(x, memo, *rv)
        ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/lib/python3.11/copy.py", line 271, in _reconstruct
    state = deepcopy(state, memo)
            ^^^^^^^^^^^^^^^^^^^^^
  File "/lib/python3.11/copy.py", line 146, in deepcopy
    y = copier(x, memo)
        ^^^^^^^^^^^^^^^
  File "/lib/python3.11/copy.py", line 231, in _deepcopy_dict
    y[deepcopy(key, memo)] = deepcopy(value, memo)
                             ^^^^^^^^^^^^^^^^^^^^^
  File "/lib/python3.11/copy.py", line 161, in deepcopy
    rv = reductor(4)
         ^^^^^^^^^^^
TypeError: cannot pickle 'pennylane_lightning.lightning_gpu_ops.DevPool' object

mlxd commented 5 months ago

Hi @TheRisenPhoenix Thanks for the report on this. The object in question is generally tied to specific local instances of GPUs on the system you are running, so we do not have serialisation rules in place for it. However, we can take a look at what seems sensible here and update you when we have a resolution.

TheRisenPhoenix commented 5 months ago

I appreciate your efforts in looking into this matter and keeping me informed! Looking forward to hearing back :)

PennyLaneAI / pennylane-lightning

[BUG] Exception when deepcopy because of non-pickle-able lightning device #677

Issue description

Source code and tracebacks