Repository: PennyLaneAI/pennylane (python)

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.

Pull request:
Transforms shallow copy qnode
Closes #4731 [sc-48599]

Currently, transforms deep copy the qnode when they are applied.  Unfortunately, this also deep copies the device.

Copying the device leads to issues like:

1) pickling errors for `lightning.kokkos` (see above issue)

2) Unnecessary allocation of large amounts of memory 

3) No longer being able to use the original devices' tracker:

```
dev = qml.device('default.qubit')

@qml.transforms.compile
@qml.qnode(dev)
def circuit():
    return qml.state()

with dev.tracker:
    circuit()
dev.tracker.totals
```
```
{}
```

This bugfix downgrades the deep copy of a qnode to a shallow copy that preserves the device.

To better separate the original from the copy, QNode now defines a custom `__copy__` method that shallow copies the `execute_kwargs`, `transform_program`, and `gradient_kwargs`. These were the three fully mutable properties of the qnode.

Patch:
diff --git a/doc/releases/changelog-0.33.0.md b/doc/releases/changelog-0.33.0.md
index 33f3de25d..e16c38b81 100644
--- a/doc/releases/changelog-0.33.0.md
+++ b/doc/releases/changelog-0.33.0.md
@@ -382,6 +382,11 @@
   decomposition.
   [(#4675)](https://github.com/PennyLaneAI/pennylane/pull/4675)
 
+* Shallow copies of the `QNode` now also copy the `execute_kwargs` and transform program. When applying
+  a transform to a `QNode`, the new qnode is only a shallow copy of the original and thus keeps the same
+  device.
+  [(#4736)](https://github.com/PennyLaneAI/pennylane/pull/4736)
+
 <h3>Breaking changes 💔</h3>
 
 * ``qml.defer_measurements`` now raises an error if a transformed circuit measures ``qml.probs``,
diff --git a/pennylane/qnode.py b/pennylane/qnode.py
index 1a9362f14..63c0d2227 100644
--- a/pennylane/qnode.py
+++ b/pennylane/qnode.py
@@ -482,6 +482,19 @@ class QNode:
         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
index 0c1004b1f..6b95e0770 100644
--- a/pennylane/transforms/core/transform_dispatcher.py
+++ b/pennylane/transforms/core/transform_dispatcher.py
@@ -196,16 +196,8 @@ class TransformDispatcher:
         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
index fc4ca54bc..08ed2b79a 100644
--- a/tests/test_qnode.py
+++ b/tests/test_qnode.py
@@ -56,6 +56,27 @@ class CustomDeviceWithDiffMethod(qml.devices.Device):
         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
index 5dbb1d6bc..2ae36088c 100644
--- a/tests/transforms/test_experimental/test_transform_dispatcher.py
+++ b/tests/transforms/test_experimental/test_transform_dispatcher.py
@@ -182,6 +182,12 @@ class TestTransformDispatcher:  # pylint: disable=too-many-public-methods
         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(