Repository: PennyLaneAI/pennylane

Issue:
[BUG] Split non commuting fails with lightning.kokkos
### Expected behavior

Split non commuting works for lightning kokkos

### Actual behavior

An exception is thrown.

### Additional information

Not an urgent bug, just looking at some tests I have for catalyst transforms.

### Source code

```shell
import jax
import pennylane as qml
import pennylane.transforms

def test_split_non_commuting():
    @qml.transforms.split_non_commuting
    @qml.qnode(qml.device("lightning.kokkos", wires=6), interface="jax")
    def interpreted():
        qml.Hadamard(1)
        qml.Hadamard(0)
        qml.PauliZ(0)
        qml.Hadamard(3)
        qml.Hadamard(5)
        qml.T(5)
        return (
            qml.expval(qml.PauliZ(0) @ qml.PauliZ(1)),
            qml.expval(qml.PauliX(0)),
            qml.expval(qml.PauliZ(1)),
            qml.expval(qml.PauliX(1) @ qml.PauliX(4)),
            qml.expval(qml.PauliX(3)),
            qml.expval(qml.PauliY(5)),
        )

    interpreted()

test_split_non_commuting()
```


### Tracebacks

```shell
python bug.py 
Kokkos::OpenMP::initialize WARNING: OMP_PROC_BIND environment variable not set
  In general, for best performance with OpenMP 4.0 or better set OMP_PROC_BIND=spread and OMP_PLACES=threads
  For best performance with OpenMP 3.1 set OMP_PROC_BIND=true
  For unit testing set OMP_PROC_BIND=false

Traceback (most recent call last):
  File "/home/erick.ochoalopez/Code/cataliist/bug.py", line 26, in <module>
    test_split_non_commuting()
  File "/home/erick.ochoalopez/Code/cataliist/bug.py", line 8, in test_split_non_commuting
    def interpreted():
  File "/home/erick.ochoalopez/Code/env/lib/python3.10/site-packages/pennylane/transforms/core/transform_dispatcher.py", line 96, in __call__
    return self._qnode_transform(obj, targs, tkwargs)
  File "/home/erick.ochoalopez/Code/env/lib/python3.10/site-packages/pennylane/transforms/core/transform_dispatcher.py", line 208, in default_qnode_transform
    qnode = copy.deepcopy(qnode)
  File "/home/erick.ochoalopez/.pyenv/versions/3.10.11/lib/python3.10/copy.py", line 172, in deepcopy
    y = _reconstruct(x, memo, *rv)
  File "/home/erick.ochoalopez/.pyenv/versions/3.10.11/lib/python3.10/copy.py", line 271, in _reconstruct
    state = deepcopy(state, memo)
  File "/home/erick.ochoalopez/.pyenv/versions/3.10.11/lib/python3.10/copy.py", line 146, in deepcopy
    y = copier(x, memo)
  File "/home/erick.ochoalopez/.pyenv/versions/3.10.11/lib/python3.10/copy.py", line 231, in _deepcopy_dict
    y[deepcopy(key, memo)] = deepcopy(value, memo)
  File "/home/erick.ochoalopez/.pyenv/versions/3.10.11/lib/python3.10/copy.py", line 172, in deepcopy
    y = _reconstruct(x, memo, *rv)
  File "/home/erick.ochoalopez/.pyenv/versions/3.10.11/lib/python3.10/copy.py", line 271, in _reconstruct
    state = deepcopy(state, memo)
  File "/home/erick.ochoalopez/.pyenv/versions/3.10.11/lib/python3.10/copy.py", line 146, in deepcopy
    y = copier(x, memo)
  File "/home/erick.ochoalopez/.pyenv/versions/3.10.11/lib/python3.10/copy.py", line 231, in _deepcopy_dict
    y[deepcopy(key, memo)] = deepcopy(value, memo)
  File "/home/erick.ochoalopez/.pyenv/versions/3.10.11/lib/python3.10/copy.py", line 161, in deepcopy
    rv = reductor(4)
TypeError: cannot pickle 'pennylane_lightning_kokkos.lightning_kokkos_qubit_ops.LightningKokkos_C128' object
```


### System information

```shell
Python 3.10.11 (main, May 13 2023, 12:07:51) [GCC 9.4.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import pennylane as qml
>>> qml.about()
Name: PennyLane
Version: 0.34.0.dev0
Summary: PennyLane is a Python quantum machine learning library by Xanadu Inc.
Home-page: https://github.com/PennyLaneAI/pennylane
Author: 
Author-email: 
License: Apache License 2.0
Location: /home/erick.ochoalopez/Code/env/lib/python3.10/site-packages
Requires: appdirs, autograd, autoray, cachetools, networkx, numpy, pennylane-lightning, requests, rustworkx, scipy, semantic-version, toml, typing-extensions
Required-by: pennylane-catalyst, PennyLane-Lightning, PennyLane-Lightning-Kokkos

Platform info:           Linux-5.15.0-84-generic-x86_64-with-glibc2.31
Python version:          3.10.11
Numpy version:           1.26.1
Scipy version:           1.11.3
Installed devices:
- default.gaussian (PennyLane-0.34.0.dev0)
- default.mixed (PennyLane-0.34.0.dev0)
- default.qubit (PennyLane-0.34.0.dev0)
- default.qubit.autograd (PennyLane-0.34.0.dev0)
- default.qubit.jax (PennyLane-0.34.0.dev0)
- default.qubit.legacy (PennyLane-0.34.0.dev0)
- default.qubit.tf (PennyLane-0.34.0.dev0)
- default.qubit.torch (PennyLane-0.34.0.dev0)
- default.qutrit (PennyLane-0.34.0.dev0)
- null.qubit (PennyLane-0.34.0.dev0)
- lightning.kokkos (PennyLane-Lightning-Kokkos-0.32.0)
- lightning.qubit (PennyLane-Lightning-0.32.0)
>>>
```


### Existing GitHub issues

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


Discussion:
The problem is copying the device. Minimal failing example:

```
dev = qml.device("lightning.kokkos", wires=6)
copy.deepcopy(dev)
```

Currently, applying a transform deep copies the entire qnode, which also performs a deep copy on the device.


Option 1:

 @trbromley @rmoyard 

When a transform is applied to the qnode, we can continue using the same device.

In addition to fixing this problem, this change would also ensure consistency with the device tracker and avoid duplicating large chunks of memory.

With larger numbers of wires, devices like `lightning.qubit` carry around a massive memory footprint.  We don't want to allocate more memory for a device copy if we don't need to.

Option 2:

Add `__copy__` and `__deepcopy__` methods to all of our devices.

Resolving patch:
diff --git a/pennylane/qnode.py b/pennylane/qnode.py
--- a/pennylane/qnode.py
+++ b/pennylane/qnode.py
@@ -482,6 +482,19 @@ def __init__(
         functools.update_wrapper(self, func)
         self._transform_program = qml.transforms.core.TransformProgram()
 
+    def __copy__(self):
+        copied_qnode = QNode.__new__(QNode)
+        for attr, value in vars(self).items():
+            if attr not in {"execute_kwargs", "_transform_program", "gradient_kwargs"}:
+                setattr(copied_qnode, attr, value)
+
+        copied_qnode.execute_kwargs = dict(self.execute_kwargs)
+        copied_qnode._transform_program = qml.transforms.core.TransformProgram(
+            self.transform_program
+        )  # pylint: disable=protected-access
+        copied_qnode.gradient_kwargs = dict(self.gradient_kwargs)
+        return copied_qnode
+
     def __repr__(self):
         """String representation."""
         if isinstance(self.device, qml.devices.Device):
diff --git a/pennylane/transforms/core/transform_dispatcher.py b/pennylane/transforms/core/transform_dispatcher.py
--- a/pennylane/transforms/core/transform_dispatcher.py
+++ b/pennylane/transforms/core/transform_dispatcher.py
@@ -196,16 +196,8 @@ def default_qnode_transform(self, qnode, targs, tkwargs):
         The default method that takes in a QNode and returns another QNode
         with the transform applied.
         """
-        # pylint: disable=protected-access
-        if (
-            isinstance(qnode._original_device, qml.Device)
-            and hasattr(qnode._original_device, "_state")
-            and qml.math.is_abstract(qnode._original_device._state)
-        ):
-            qnode._original_device.reset()
-            qnode.device.reset()
-
-        qnode = copy.deepcopy(qnode)
+
+        qnode = copy.copy(qnode)
 
         if self.expand_transform:
             qnode.add_transform(TransformContainer(self._expand_transform, targs, tkwargs))


Test patch:
diff --git a/tests/test_qnode.py b/tests/test_qnode.py
--- a/tests/test_qnode.py
+++ b/tests/test_qnode.py
@@ -56,6 +56,27 @@ def compute_derivatives(self, circuits, execution_config=None):
         return 0
 
 
+def test_copy():
+    """Test that a shallow copy also copies the execute kwargs, gradient kwargs, and transform program."""
+    dev = CustomDevice()
+
+    qn = qml.QNode(dummyfunc, dev)
+    copied_qn = copy.copy(qn)
+    assert copied_qn is not qn
+    assert copied_qn.execute_kwargs == qn.execute_kwargs
+    assert copied_qn.execute_kwargs is not qn.execute_kwargs
+    assert list(copied_qn.transform_program) == list(qn.transform_program)
+    assert copied_qn.transform_program is not qn.transform_program
+    assert copied_qn.gradient_kwargs == qn.gradient_kwargs
+    assert copied_qn.gradient_kwargs is not qn.gradient_kwargs
+
+    assert copied_qn.func is qn.func
+    assert copied_qn.device is qn.device
+    assert copied_qn.interface is qn.interface
+    assert copied_qn.diff_method == qn.diff_method
+    assert copied_qn.expansion_strategy == qn.expansion_strategy
+
+
 # pylint: disable=too-many-public-methods
 class TestValidation:
     """Tests for QNode creation and validation"""
diff --git a/tests/transforms/test_experimental/test_transform_dispatcher.py b/tests/transforms/test_experimental/test_transform_dispatcher.py
--- a/tests/transforms/test_experimental/test_transform_dispatcher.py
+++ b/tests/transforms/test_experimental/test_transform_dispatcher.py
@@ -182,6 +182,12 @@ def qnode_circuit(a):
         qnode_transformed = dispatched_transform(qnode_circuit, 0)
         assert not qnode_circuit.transform_program
 
+        assert qnode_transformed.device is qnode_circuit.device
+
+        with dev.tracker:
+            qnode_circuit(0.1)
+        assert dev.tracker.totals["executions"] == 1
+
         assert isinstance(qnode_transformed, qml.QNode)
         assert isinstance(qnode_transformed.transform_program, qml.transforms.core.TransformProgram)
         assert isinstance(