Repository: PennyLaneAI/pennylane

Issue:
[BUG] Loading hybrid quantum neuronal networks does not work with TensorFlow
### Expected behavior

When using the tensorflow function `model.save` und `tf.keras.models.load_model` it is expected that the model load all layers and its ways and recreates the saved model to include all weights etc.


### Actual behavior

When using  `model.save` and `tf.keras.models.load_model` with a hybrid model including both classical layers and `qml.qnn.KerasLayerl` the network including the quantum layer is reloaded without error, however the weights of the quantum layer are not. 
This issue exists independent of the order of quantum and classical layers.

### Additional information

This issue was discussed in the pennylane support threat 
https://discuss.pennylane.ai/t/error-reloading-circuit-from-qasm-string/3679
This threat already gave rise to issue https://github.com/PennyLaneAI/pennylane/issues/4856 to which this issue is possibly related

### Source code

```shell
import tensorflow as tf
import pennylane as qml
import numpy as np

dev = qml.device("default.qubit")
n_qubits = 2

def create_model():
    @qml.qnode(dev)
    def circuit(inputs, weights):
        qml.AngleEmbedding(inputs, wires=range(n_qubits))
        qml.RX(weights[0], 0)
        qml.RX(weights[1], 1)
        return [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))]

    weight_shapes = {"weights": (n_qubits,)}
    quantum_layer = qml.qnn.KerasLayer(circuit, weight_shapes, output_dim=2)

    model = tf.keras.models.Sequential()
    model.add(quantum_layer)
    model.add(tf.keras.layers.Dense(2, activation=tf.nn.softmax, input_shape=(2,)))

    model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])
    return model

model = create_model()

num_points = 5
dummy_input_data = np.random.uniform(0, np.pi, size=(num_points, 2))
dummy_output_data = np.random.randint(2, size=(num_points, 2))

model.fit(dummy_input_data, dummy_output_data, epochs=1, batch_size=0)

model.save("model")

loaded_model = tf.keras.models.load_model("model")

print("")
print("saved weights:", model.layers[0].weights)
print("")
print("Loaded weights:", loaded_model.layers[0].weights)
```


### Tracebacks

```shell
WARNING:tensorflow:You are casting an input of type complex128 to an incompatible dtype float32.  This will discard the imaginary part and may not be what you intended.
WARNING:tensorflow:You are casting an input of type complex128 to an incompatible dtype float32.  This will discard the imaginary part and may not be what you intended.
WARNING:tensorflow:You are casting an input of type complex128 to an incompatible dtype float32.  This will discard the imaginary part and may not be what you intended.
WARNING:tensorflow:You are casting an input of type complex128 to an incompatible dtype float32.  This will discard the imaginary part and may not be what you intended.
1/1 [==============================] - 0s 349ms/step - loss: 0.5533 - accuracy: 0.8000
WARNING:tensorflow:AutoGraph could not transform <function validate_device_wires at 0x0000027642280F70> and will run it as-is.
Please report this to the TensorFlow team. When filing the bug, set the verbosity to 10 (on Linux, `export AUTOGRAPH_VERBOSITY=10`) and attach the full output.
Cause: 'NoneType' object has no attribute '_fields'
To silence this warning, decorate the function with @tf.autograph.experimental.do_not_convert
WARNING: AutoGraph could not transform <function validate_device_wires at 0x0000027642280F70> and will run it as-is.
Please report this to the TensorFlow team. When filing the bug, set the verbosity to 10 (on Linux, `export AUTOGRAPH_VERBOSITY=10`) and attach the full output.
Cause: 'NoneType' object has no attribute '_fields'
To silence this warning, decorate the function with @tf.autograph.experimental.do_not_convert
WARNING:tensorflow:AutoGraph could not transform <function _gcd_import at 0x000002762FE83490> and will run it as-is.
Cause: Unable to locate the source code of <function _gcd_import at 0x000002762FE83490>. Note that functions defined in certain environments, like the interactive Python shell, do not expose their source code. If that is the case, you should define them in a .py source file. If you are certain the code is graph-compatible, wrap the call using @tf.autograph.experimental.do_not_convert. Original error: could not get source code
To silence this warning, decorate the function with @tf.autograph.experimental.do_not_convert
WARNING: AutoGraph could not transform <function _gcd_import at 0x000002762FE83490> and will run it as-is.
Cause: Unable to locate the source code of <function _gcd_import at 0x000002762FE83490>. Note that functions defined in certain environments, like the interactive Python shell, do not expose their source code. If that is the case, you should define them in a .py source file. If you are certain the code is graph-compatible, wrap the call using @tf.autograph.experimental.do_not_convert. Original error: could not get source code
To silence this warning, decorate the function with @tf.autograph.experimental.do_not_convert
INFO:tensorflow:Assets written to: model\assets
INFO:tensorflow:Assets written to: model\assets

saved weights: [<tf.Variable 'weights:0' shape=(2,) dtype=float32, numpy=array([ 0.49809802, -0.9938641 ], dtype=float32)>]

Loaded weights: []
```


### System information

```shell
Platform info:           Windows-10-10.0.19045-SP0
Python version:          3.10.12
Numpy version:           1.26.0
Scipy version:           1.11.2
Installed devices:
- default.gaussian (PennyLane-0.33.1)
- default.mixed (PennyLane-0.33.1)
- default.qubit (PennyLane-0.33.1)
- default.qubit.autograd (PennyLane-0.33.1)
- default.qubit.jax (PennyLane-0.33.1)
- default.qubit.legacy (PennyLane-0.33.1)
- default.qubit.tf (PennyLane-0.33.1)
- default.qubit.torch (PennyLane-0.33.1)
- default.qutrit (PennyLane-0.33.1)
- null.qubit (PennyLane-0.33.1)
- lightning.qubit (PennyLane-Lightning-0.33.1)
- strawberryfields.fock (PennyLane-SF-0.29.0)
- strawberryfields.gaussian (PennyLane-SF-0.29.0)
- strawberryfields.gbs (PennyLane-SF-0.29.0)
- strawberryfields.remote (PennyLane-SF-0.29.0)
- strawberryfields.tf (PennyLane-SF-0.29.0)
- qiskit.aer (PennyLane-qiskit-0.32.0)
- qiskit.basicaer (PennyLane-qiskit-0.32.0)
- qiskit.ibmq (PennyLane-qiskit-0.32.0)
- qiskit.ibmq.circuit_runner (PennyLane-qiskit-0.32.0)
- qiskit.ibmq.sampler (PennyLane-qiskit-0.32.0)
- qiskit.remote (PennyLane-qiskit-0.32.0)

The issue is also visible in the Tracebacks displayed in https://github.com/PennyLaneAI/pennylane/issues/4856 and therefore also reproducible the system configuration displayed there.
```


### Existing GitHub issues

- [X] I have searched existing GitHub issues to make sure the issue does not already exist.
[BUG] qubit device expval method doesn't work with batch and BasisStateProjector
### Expected behavior

When utilizing a qubit device with batch processing, expval output, and BasisStateProjector, the probabilities need to be correctly indexed in the good dimension.

### Actual behavior

When utilizing a qubit device with batch processing, expval output, and BasisStateProjector, the probabilities are incorrectly index in the wrong dimension. This result in the expval output returning the probabilities of a vector within the batch, rather than providing the probabilities of the desired state for all entries in the batch.

### Additional information

A PR is already open for this issue. See https://github.com/PennyLaneAI/pennylane/pull/4971 .

### Source code

```shell
import pennylane as qml
from pennylane.devices import DefaultQubit
from pennylane.ops.qubit.observables import BasisStateProjector
import numpy as np


def circuit(x, y):
    n_wires = qml.math.shape(x)[-1]
    for i in range(n_wires):
        qml.RX(x[:, i], wires=i)
        qml.RY(y[:, i], wires=i)
    projector: BasisStateProjector = qml.Projector(np.zeros(n_wires), wires=range(n_wires))
    return qml.expval(projector)


if __name__ == '__main__':
    device = DefaultQubit(wires=2)
    qnode = qml.QNode(circuit, device)
    batch_size = device.num_wires * 3 + 1
    x = np.random.uniform(0, 2 * np.pi, size=(batch_size, device.num_wires))
    y = np.random.uniform(0, 2 * np.pi, size=(batch_size, device.num_wires))
    out = qnode(x, y)
    np.testing.assert_allclose(qml.math.shape(out), (batch_size,))
```


### Tracebacks

```shell
Traceback (most recent call last):
  File "E:\Github\FermionicSimulation\scratches\pennylane_pr.py", line 23, in <module>
    np.testing.assert_allclose(qml.math.shape(out), (batch_size,))
  File "E:\Github\FermionicSimulation\venv\lib\site-packages\numpy\testing\_private\utils.py", line 1527, in assert_allclose
    assert_array_compare(compare, actual, desired, err_msg=str(err_msg),
  File "E:\Github\FermionicSimulation\venv\lib\site-packages\numpy\testing\_private\utils.py", line 844, in assert_array_compare
    raise AssertionError(msg)
AssertionError: 
Not equal to tolerance rtol=1e-07, atol=0

Mismatched elements: 1 / 1 (100%)
Max absolute difference: 3
Max relative difference: 0.42857143
 x: array([4])
 y: array([7])
```


### System information

```shell
Name: PennyLane
Version: 0.32.0
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: e:\github\fermionicsimulation\venv\lib\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:           Windows-10-10.0.22621-SP0
Python version:          3.9.12
Numpy version:           1.23.5
Scipy version:           1.11.3
Installed devices:
- default.gaussian (PennyLane-0.32.0)
- default.mixed (PennyLane-0.32.0)
- default.qubit (PennyLane-0.32.0)
- default.qubit.autograd (PennyLane-0.32.0)
- default.qubit.jax (PennyLane-0.32.0)
- default.qubit.tf (PennyLane-0.32.0)
- default.qubit.torch (PennyLane-0.32.0)
- default.qutrit (PennyLane-0.32.0)
- null.qubit (PennyLane-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:

Resolving patch:
diff --git a/pennylane/wires.py b/pennylane/wires.py
--- a/pennylane/wires.py
+++ b/pennylane/wires.py
@@ -15,10 +15,10 @@
 This module contains the :class:`Wires` class, which takes care of wire bookkeeping.
 """
 import functools
-from collections.abc import Iterable, Sequence
 import itertools
-
+from collections.abc import Iterable, Sequence
 import numpy as np
+from pennylane.pytrees import register_pytree
 
 
 class WireError(Exception):
@@ -99,6 +99,15 @@ class Wires(Sequence):
          wires (Any): the wire label(s)
     """
 
+    def _flatten(self):
+        """Serialize Wires into a flattened representation according to the PyTree convension."""
+        return self._labels, ()
+
+    @classmethod
+    def _unflatten(cls, data, _metadata):
+        """De-serialize flattened representation back into the Wires object."""
+        return cls(data, _override=True)
+
     def __init__(self, wires, _override=False):
         if _override:
             self._labels = wires
@@ -481,3 +490,7 @@ def unique_wires(list_of_wires):
                     unique.append(wire)
 
         return Wires(tuple(unique), _override=True)
+
+
+# Register Wires as a PyTree-serializable class
+register_pytree(Wires, Wires._flatten, Wires._unflatten)  # pylint: disable=protected-access
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 """Setup file for package installation."""
+# pylint: disable=unspecified-encoding, consider-using-with
 
 from setuptools import setup, find_packages
 
@@ -30,7 +31,7 @@
     "semantic-version>=2.7",
     "autoray>=0.6.1",
     "cachetools",
-    "pennylane-lightning>=0.33",
+    "pennylane-lightning>=0.34",
     "requests",
     "typing_extensions",
 ]


Test patch:
diff --git a/tests/test_wires.py b/tests/test_wires.py
--- a/tests/test_wires.py
+++ b/tests/test_wires.py
@@ -362,3 +362,17 @@ def test_hash_cached(self):
         assert wires._hash is None
         h = hash(wires)
         assert wires._hash == h
+
+    @pytest.mark.jax
+    @pytest.mark.parametrize(
+        "source", [1, -2, "a", "q1", -1.4, np.array([0, 1, 2]), [0, 1, 2], (0, 1, 2), range(3)]
+    )
+    def test_wires_pytree(self, source):
+        """Test that Wires class supports the PyTree flattening interface"""
+        from jax.tree_util import tree_flatten, tree_unflatten
+
+        wires = Wires(source)
+        wires_flat, tree = tree_flatten(wires)
+        wires2 = tree_unflatten(tree, wires_flat)
+        assert isinstance(wires2, Wires), f"{wires2} is not Wires"
+        assert wires == wires2, f"{wires} != {wires2}"