[BUG] Loading and saving models and weights with TensorFlow doesn't behave as expected

[isaacdevlugt]

Expected behavior

When save_weights and load_weights are called, I expect that the model's weights are the same before and after they're saved and loaded. Also, if I use model.save to save the whole model and tf.keras.models.load_model, I expect that the weights are the same.

Actual behavior

In both cases (saving just the weights or saving the whole model), the models that are created / loaded after saving give the same outputs / predictions, but when their weights are inspected, what's shown isn't the same.

Additional information

https://discuss.pennylane.ai/t/error-reloading-circuit-from-qasm-string/3679

Source code

import tensorflow as tf
import pennylane as qml
from pennylane 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("saved weights:", model.weights)
print("Loaded weights:", loaded_model.weights)

print(model.predict(dummy_input_data) == loaded_model.predict(dummy_input_data))

saved_weights = model.weights
model.save_weights("saved_weights")

new_model = create_model()
new_model.load_weights("saved_weights")
print("Weights saved:", saved_weights)
print()
print("Loaded weights:", new_model.weights)
print()
print(model.predict(dummy_input_data) == new_model.predict(dummy_input_data))

Tracebacks

saved weights: [<tf.Variable 'weights:0' shape=(2,) dtype=float32, numpy=array([-0.6488751,  1.1488918], dtype=float32)>, <tf.Variable 'dense/kernel:0' shape=(2, 2) dtype=float32, numpy=
array([[ 0.1556892, -0.646571 ],
       [-0.2932235,  1.1476017]], dtype=float32)>, <tf.Variable 'dense/bias:0' shape=(2,) dtype=float32, numpy=array([-0.00099998,  0.00099998], dtype=float32)>]
Loaded weights: [<tf.Variable 'dense/kernel:0' shape=(2, 2) dtype=float32, numpy=
array([[ 0.1556892, -0.646571 ],
       [-0.2932235,  1.1476017]], dtype=float32)>, <tf.Variable 'dense/bias:0' shape=(2,) dtype=float32, numpy=array([-0.00099998,  0.00099998], dtype=float32)>]
1/1 [==============================] - 0s 30ms/step
2023-11-17 16:19:22.491641: I tensorflow/core/common_runtime/executor.cc:1197] [/device:CPU:0] (DEBUG INFO) Executor start aborting (this does not indicate an error and you can ignore this message): INVALID_ARGUMENT: You must feed a value for placeholder tensor 'inputs' with dtype float and shape [?,2]
     [[{{node inputs}}]]
2023-11-17 16:19:22.502508: I tensorflow/core/common_runtime/executor.cc:1197] [/device:CPU:0] (DEBUG INFO) Executor start aborting (this does not indicate an error and you can ignore this message): INVALID_ARGUMENT: You must feed a value for placeholder tensor 'Placeholder' with dtype float and shape [?,2]
     [[{{node Placeholder}}]]
1/1 [==============================] - 0s 138ms/step
[[ True  True]
 [ True  True]
 [ True  True]
 [ True  True]
 [ True  True]]

Weights saved: [<tf.Variable 'weights:0' shape=(2,) dtype=float32, numpy=array([-0.6488751,  1.1488918], dtype=float32)>, <tf.Variable 'dense/kernel:0' shape=(2, 2) dtype=float32, numpy=
array([[ 0.1556892, -0.646571 ],
       [-0.2932235,  1.1476017]], dtype=float32)>, <tf.Variable 'dense/bias:0' shape=(2,) dtype=float32, numpy=array([-0.00099998,  0.00099998], dtype=float32)>]

Loaded weights: [<tf.Variable 'weights:0' shape=(2,) dtype=float32, numpy=array([-0.6488751,  1.1488918], dtype=float32)>]

1/1 [==============================] - 0s 31ms/step
1/1 [==============================] - 0s 39ms/step
[[ True  True]
 [ True  True]
 [ True  True]
 [ True  True]
 [ True  True]]

System information

Name: PennyLane
Version: 0.33.1
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: /Users/isaac/.virtualenvs/pennylane-tensorflow/lib/python3.9/site-packages
Requires: appdirs, autograd, autoray, cachetools, networkx, numpy, pennylane-lightning, requests, rustworkx, scipy, semantic-version, toml, typing-extensions
Required-by: PennyLane-Lightning, PennyLane-qiskit

Platform info:           macOS-14.1.1-x86_64-i386-64bit
Python version:          3.9.14
Numpy version:           1.23.5
Scipy version:           1.11.3
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)
...
- qiskit.ibmq.circuit_runner (PennyLane-qiskit-0.33.0)
- qiskit.ibmq.sampler (PennyLane-qiskit-0.33.0)
- qiskit.remote (PennyLane-qiskit-0.33.0)
- lightning.qubit (PennyLane-Lightning-0.33.1)

Existing GitHub issues

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

[trbromley]

Thanks @isaacdevlugt! This is interesting, because we do test for model weight saving here: https://github.com/PennyLaneAI/pennylane/blob/master/tests/qnn/test_keras.py#L578

However, your example above places a classical layer after the quantum one, while our test just has a classical layer before.

[lillian542]

The user was able to solve the problem by updating how AngleEmbedding was initialized in their code, so we are closing this issue. The model.save function uses autograph, which implicitly adds a batch dimension.