How to get density matrix before measuring the circuit? Snapshot method does not work with Tensorflow

Hi! I am trying to get a snapshot of the density matrix before measuring a circuit run in a Hybrid CNN with tensorflow. However I am not able to achieve a good code behaviour. I used the fuctional Keras API to make the CNN work, i.e., I created a new class for the model where I defined the convolution circuit as:

# Create a quantum device
self.dev = qml.device(device, wires=n_qubits)

@tf.function(jit_compile=None)
@qml.qnode(self.dev)
def circuit(inputs, params):
    # Apply trainable encoding operations on each qubit
    for i in range(self.n_qubits):
        qml.Rot(params[i] * inputs[i] + params[i + self.n_qubits],
                params[i + 2 * self.n_qubits] * inputs[i] + params[i + 3 * self.n_qubits],
                params[i + 4 * self.n_qubits] * inputs[i] + params[i + 5 * self.n_qubits],
                wires=i)

    # Apply circular entanglement
    for i in range(self.n_qubits):
        qml.CNOT(wires=[i, (i + 1) % self.n_qubits])

    # Save state vector
    qml.Snapshot()

    # Measure all qubits
    return [qml.expval(qml.PauliZ(i)) for i in range(self.n_qubits)]

self.circuit = circuit

Then I try to use the circuit here and get the state vector with qml.snapshots:

for i in range(batch_size):
    for j in range(0, self.image_height, self.stride):
        for k in range(0, self.image_width, self.stride):

            # Process a squared 2x2 region of the image with a quantum circuit
            patch = tf.reshape(inputs[i, j:j+2, k:k+2, 0], [4])
            result_dict = qml.snapshots(self.circuit)(patch, self.params)
            measurement_expvals = result_dict['execution_results']
            out[i][j // self.stride][k // self.stride] = measurement_expvals

            # Get entanglement entropies from reduced density matrices
            # density_matrix = qml.math.dm_from_state_vector(result_dict[0])
            # qubits_idxs = list(range(self.n_qubits))
            # for l in range(1, self.n_qubits//2+1):
            #     for m in itertools.combinations(qubits_idxs, l):
            #         qubits_to_keep = [q for q in qubits_idxs if q not in m]
            #         entanglement_entropies.append(qml.math.vn_entropy(density_matrix, qubits_to_keep, base=2))

However, when calling the model with some data, I get this error message:

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Cell In[18], line 3
      2 with tf.device('/CPU:0'):
----> 3     qccnn.call(tf.convert_to_tensor(train_images[:batch_size]))

File ~/.virtualenvs/qml/lib/python3.10/site-packages/keras/engine/sequential.py:425, in Sequential.call(self, inputs, training, mask)
    422 if "training" in argspec:
    423     kwargs["training"] = training
--> 425 outputs = layer(inputs, **kwargs)
    427 if len(tf.nest.flatten(outputs)) != 1:
    428     raise ValueError(SINGLE_LAYER_OUTPUT_ERROR_MSG)

File ~/.virtualenvs/qml/lib/python3.10/site-packages/keras/utils/traceback_utils.py:70, in filter_traceback.<locals>.error_handler(*args, **kwargs)
     67     filtered_tb = _process_traceback_frames(e.__traceback__)
     68     # To get the full stack trace, call:
     69     # `tf.debugging.disable_traceback_filtering()`
---> 70     raise e.with_traceback(filtered_tb) from None
     71 finally:
     72     del filtered_tb

Cell In[15], line 111, in QuantumConv.call(self, inputs, **kwargs)
    108                     patch = tf.reshape(inputs[i, j:j+2, k:k+2, 0], [4])
--> 111                     result_dict = qml.snapshots(self.circuit)(patch, self.params)
    112                     measurement_expvals = result_dict['execution_results']

File ~/.virtualenvs/qml/lib/python3.10/site-packages/pennylane/debugging.py:92, in snapshots.<locals>.get_snapshots(*args, **kwargs)
     91 def get_snapshots(*args, **kwargs):
---> 92     old_interface = qnode.interface
     93     if old_interface == "auto":
     94         qnode.interface = qml.math.get_interface(*args, *list(kwargs.values()))

AttributeError: Exception encountered when calling layer "quantum_conv" "                 f"(type QuantumConv).

'Function' object has no attribute 'interface'

Call arguments received by layer "quantum_conv" "                 f"(type QuantumConv):
  • inputs=tf.Tensor(shape=(8, 18, 12, 1), dtype=float32)
  • kwargs={'training': 'None'}

If I remove the decorator @tf.function(jit_compile=None), then the code does not throw this error but it runs forever and never stops. OR if I don’t use the qml.snapshots method then it does not throw any error.

Am I using an outdated method? How do I get the state density matrix before measuring the circuit?


Output of qml.about():

Name: PennyLane
Version: 0.35.0
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: /qcfs/bravo/.virtualenvs/qml/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-Cirq, PennyLane-qiskit, PennyLane_Lightning, PennyLane_Lightning_GPU

Platform info:           Linux-5.15.0-92-generic-x86_64-with-glibc2.35
Python version:          3.10.12
Numpy version:           1.23.5
Scipy version:           1.11.3
Installed devices:
- cirq.mixedsimulator (PennyLane-Cirq-0.34.0)
- cirq.pasqal (PennyLane-Cirq-0.34.0)
- cirq.qsim (PennyLane-Cirq-0.34.0)
- cirq.qsimh (PennyLane-Cirq-0.34.0)
- cirq.simulator (PennyLane-Cirq-0.34.0)
- lightning.qubit (PennyLane_Lightning-0.35.0)
- default.clifford (PennyLane-0.35.0)
- default.gaussian (PennyLane-0.35.0)
- default.mixed (PennyLane-0.35.0)
- default.qubit (PennyLane-0.35.0)
- default.qubit.autograd (PennyLane-0.35.0)
- default.qubit.jax (PennyLane-0.35.0)
- default.qubit.legacy (PennyLane-0.35.0)
- default.qubit.tf (PennyLane-0.35.0)
- default.qubit.torch (PennyLane-0.35.0)
- default.qutrit (PennyLane-0.35.0)
- null.qubit (PennyLane-0.35.0)
- lightning.gpu (PennyLane_Lightning_GPU-0.35.0)
- qiskit.aer (PennyLane-qiskit-0.35.1)
- qiskit.basicaer (PennyLane-qiskit-0.35.1)
- qiskit.ibmq (PennyLane-qiskit-0.35.1)
- qiskit.ibmq.circuit_runner (PennyLane-qiskit-0.35.1)
- qiskit.ibmq.sampler (PennyLane-qiskit-0.35.1)
- qiskit.remote (PennyLane-qiskit-0.35.1)

Hey @joaofbravo,

Parts of your code are missing and I think I need more context to understand what you’re trying to do. Are you using KerasLayer (qml.qnn.KerasLayer — PennyLane 0.35.1 documentation) to inject this into your model? Is [qml.expval(qml.PauliZ(i)) for i in range(self.n_qubits)] what you want to use as the quantum layer’s output that feeds into another layer as input? And you’re just wanting to use the state vector / density matrix at the end of the quantum layer for something else that happens during the forward pass?

If you could shrink your code down into a working toy example that exemplifies what you’re trying to accomplish, that would be a huge plus :slight_smile:.

Hi Isaac,

I did not use KerasLayer (qml.qnn.KerasLayer — PennyLane 0.35.1 documentation) because I kept getting errors with it due to how the quantum convolution is performed, so I just created a custom Layer class for it.

[qml.expval(qml.PauliZ(i)) for i in range(self.n_qubits)] is the output of the circuit which is then rearranged with out[i][j // self.stride][k // self.stride] = measurement_expvals to be the output of the Layer.

I need the state vector or density matrix of the state before measuring because I want to assess how the entanglement entropy changes along the training. So I need it for a parallel computation which has nothing to do with training.

Here is some code:
qccnn_trainable_entanglement_dummy.py (9.4 KB)

Hey @joaofbravo,

I don’t see why you can’t use KerasLayer here :thinking:. You mentioned this:

… because I kept getting errors with it due to how the quantum convolution is performed …

I’m curious what the error was there when you tried to implement want you wanted with KerasLayer. In any case, I think at the end of the day, you are trying to accomplish this (at a high level):

  1. Have a variational circuit as part of your hybrid CNN.
  2. Have a custom forward pass that not only requires the output of the circuit, but also its density matrix before the measurement.

You can accomplish that with using KerasLayer! Here’s a very simple toy example:

import pennylane as qml
import tensorflow as tf

# create a quantum circuit that will be turned into a Keras layer

dev = qml.device("default.mixed", wires=n_qubits)

@qml.qnode(dev)
def circuit(inputs, weights): # the "inputs" arg is needed, but I won't use it for the purpose of this example
    qml.RX(weights, wires=0)

    qml.Snapshot()

    return qml.expval(qml.PauliZ(0))

weight_shapes = {"weights": (1,)}
qlayer = qml.qnn.KerasLayer(circuit, weight_shapes, output_dim=1)

Now, to access the density matrix with snapshots, you need to provide the weights:

qlayer_weights = qlayer.get_weights()[0]

Create a model:

clayer = tf.keras.layers.Dense(2)
model = tf.keras.models.Sequential([qlayer, clayer])
model.compile()

Call snapshots

qml.snapshots(circuit)(None, qlayer_weights) # NB: inputs = None
{0: array([[0.82739791+0.j        , 0.        +0.37790291j],
        [0.        -0.37790291j, 0.17260209+0.j        ]]),
 'execution_results': array(0.65479582)}

I know this is a very simple example, but you should be able to contain these concepts within your QuantumConv class (that includes a custom forward pass), but now you’re using KerasLayer and there shouldn’t (:crossed_fingers:) be any issues with PennyLane talking to Tensorflow.

Let me know if this helps!

Just to follow up, it would look something like this:

import pennylane as qml
import tensorflow as tf
from tensorflow import keras

n_qubits=2
num_test_inputs = 10

test_inputs = tf.random.uniform((num_test_inputs, 2**n_qubits,), 0, 1)

class Model(keras.layers.Layer):
    def __init__(self, n_qubits=n_qubits, name='my_model', **kwargs): 
        super().__init__(name=name, **kwargs) 
        self.n_qubits = n_qubits

        self.dev = qml.device("default.mixed", wires=n_qubits)

        self.circuit = self.create_circuit(n_qubits)

        weight_shapes = {"weights": (2,)}
        self.qlayer = qml.qnn.KerasLayer(self.circuit, weight_shapes, output_dim=2)

        self.clayer = tf.keras.layers.Dense(2)

        self.model = tf.keras.models.Sequential([self.qlayer, self.clayer])
        self.model.compile()

    def create_circuit(self, n_qubits):

        @qml.qnode(self.dev)
        def circuit(inputs, weights):
            qml.AmplitudeEmbedding(inputs, wires=range(n_qubits), normalize=True)
            qml.RX(weights[0], wires=0)
            qml.RZ(weights[1], wires=1)

            qml.Snapshot('density_matrix')

            return qml.expval(qml.Z(0)), qml.expval(qml.Z(1))
        
        return circuit
    
    def call(self, inputs):
        x = self.qlayer(inputs)
        
        # do some arbitrary stuff with density matrix
        density_matrix = qml.snapshots(self.circuit)(inputs, self.qlayer.get_weights()[0])['density_matrix']
        purity = qml.math.purity(density_matrix, list(range(self.n_qubits)))
        x *= purity
        
        x = self.clayer(x)

        return x


model = Model()
model(test_inputs)
<tf.Tensor: shape=(10, 2), dtype=float32, numpy=
array([[-0.10869376,  0.02593517],
       [-0.01298368,  0.04934663],
       [-0.06047328, -0.11851539],
       [-0.13707723,  0.15470251],
       [-0.02513164,  0.10013923],
       [-0.00613316, -0.12370723],
       [ 0.01306366,  0.05771445],
       [-0.10260486, -0.01226309],
       [ 0.15258303, -0.30820817],
       [-0.23772222,  0.02938864]], dtype=float32)>

Hi. I am not sure how to include that in my code with the qml.snapshots method. Should I make the circuit(inputs, params) a KerasLayer and call KerasLayer qml.snapshots(qlayer)? This does not work. It gives the error

TypeError: Exception encountered when calling layer “quantum_conv” " f"(type QuantumConv).

KerasLayer.call() takes 2 positional arguments but 3 were given

Call arguments received by layer “quantum_conv” " f"(type QuantumConv):
• inputs=tf.Tensor(shape=(2, 18, 12, 1), dtype=float32)
• kwargs={‘training’: ‘None’}

If I still call qml.snapshots(circuit) then the KerasLayer is unused.

The problem is that my circuit is not a layer. It is a small part of the convolutional layer. It should apply that circuit many times in parallel to patches of the image. So I cannot simply put the circuit in a tf.keras.models.Sequential.

Edit: checking your new post, it seems to be what I am looking for. I will give it a try and tell you if it worked

qml.snapshots needs to see a QNode as input, not a Keras layer (see here: qml.snapshots — PennyLane 0.35.1 documentation).

If I still call qml.snapshots(circuit) then the KerasLayer is unused

Yep! The whole point of my examples was to show that you have to go around KerasLayer to use snapshots :slight_smile:

So I cannot simply put the circuit in a tf.keras.models.Sequential

That should be fine! You don’t have to define your model that way. You can simply call on the KerasLayer object as many times as you need within your call method :slight_smile:.

1 Like

Got it! However, when I ran the example for one epoch, I get this warning now:

Epoch 1/2
WARNING:tensorflow:Gradients do not exist for variables [‘sequential/params:0’] when minimizing the loss. If you’re using model.compile(), did you forget to provide a loss argument?
1/1 - 11s - loss: 0.8179 - accuracy: 0.5000 - val_loss: 1.2412 - val_accuracy: 0.5000 - 11s/epoch - 11s/step

Also, does this mean that I need to run the circuit twice to get the density matrix (once with the KerasLayer and once with snapshots? Is there a way to merge this for efficiency?

qccnn_trainable_entanglement_dummy.py (9.4 KB)

Another issue is that this seems to work whenever the call method is called. Is there a way to store the density matrices of the data used in training and validation separately?

Apologies! I’m not a tensorflow user myself, so I had to do a bit of reading. I was misusing the Layer class from Keras (what the class above was inheriting from). This shouldn’t contain any model creation, etc. It should only contain things that are needed to define a layer.

I adapted my code to our KerasLayer tutorial: https://pennylane.ai/qml/demos/tutorial_qnn_module_tf/:slight_smile:

import pennylane as qml
import tensorflow as tf
from tensorflow import keras

import numpy as np
from sklearn.datasets import make_moons

# Set random seeds
np.random.seed(42)
tf.random.set_seed(42)

X, y = make_moons(n_samples=200, noise=0.1)
y_hot = tf.keras.utils.to_categorical(y, num_classes=2)  # one-hot encoded labels

n_qubits=2

class MyLayer(keras.layers.Layer):
    def __init__(self, n_qubits=n_qubits, name='my_layer', **kwargs): 
        super().__init__(name=name, **kwargs) 
        self.n_qubits = n_qubits

        self.dev = qml.device("default.mixed", wires=n_qubits)

        self.circuit = self.create_circuit(n_qubits)

        weight_shapes = {"weights": (2,)}
        self.qlayer = qml.qnn.KerasLayer(self.circuit, weight_shapes, output_dim=2)

        self.clayer1 = tf.keras.layers.Dense(2)
        self.clayer2 = tf.keras.layers.Dense(2, activation="softmax")
        
    def create_circuit(self, n_qubits):

        @qml.qnode(self.dev)
        def circuit(inputs, weights):
            qml.AngleEmbedding(inputs, wires=range(n_qubits))
            qml.RX(weights[0], wires=0)
            qml.RZ(weights[1], wires=1)

            qml.Snapshot('density_matrix')

            return qml.expval(qml.Z(0)), qml.expval(qml.Z(1))
        
        return circuit
    
    def call(self, inputs):
        x = self.qlayer(inputs)
        
        # do some arbitrary stuff with density matrix
        density_matrix = qml.snapshots(self.circuit)(inputs, self.qlayer.get_weights()[0])['density_matrix']
        purity = qml.math.purity(density_matrix, list(range(self.n_qubits)))
        x *= purity

        x = self.clayer1(x)
        x = self.clayer2(x)

        return x

The MyLayer class now behaves similarly to how Dense would, for example. Now you create a model using that. You can do this with Sequential:

my_layer = MyLayer()
model = tf.keras.models.Sequential([my_layer])

Now you can compile and fit:

opt = tf.keras.optimizers.SGD(learning_rate=0.2)
model.compile(opt, loss="mae", metrics=["accuracy"])

fitting = model.fit(X, y_hot, epochs=6, batch_size=5, validation_split=0.25, verbose=2)

I found out my issue. I added weights to the layer before using the KerasLayer method with self.params = self.add_weight(shape=(self.n_qubits*3*2,), initializer=keras.initializers.RandomUniform(minval=0., maxval=2*np.pi, seed=self.seed)). Since KerasLayers generates its own weights, these weights were not used and then I got the warning:

WARNING:tensorflow:Gradients do not exist for variables [‘sequential/params:0’] when minimizing the loss

Okay so it is working, but I am still getting an unwanted amount of density matrices. As I stated before, the density matrix is generated every time the call method is called. This happens both for train data and validation data every epoch. Is there a way to store the density matrices of train and validation data separately?

Is there a way to store the density matrices of train and validation data separately?

Won’t they always be different anyway, since the data that you input into the model is different in the training and validation data? Maybe I’m missing something with what you’re trying to do. In any case, do you want to store both density matrices external from your class, or inside of it?

Yes states generated by all data points will be different. The problem is that train data points and validation data points are both used during training, so if I want to track the average entanglement of the density matrices created by the training data, it does not seem so easy. After for example 100 epochs, all I get is an array with more density matrices than I want.

Storing the density matrices inside the class in two different variables would be optimal but doing it externally is also okay.

Storing the density matrices inside the class in two different variables would be optimal

Nice! That is very doable with your class. You can give your class two different density matrix attributes, update them when you make a forward pass with your training and validation, and voila!

Ah right. I was using the fit method but I guess I need to write the training procedure myself for this.

Modifying the fit method might be one way, but I think that it might be doable without doing that :thinking:. You’re more familiar with what you need accomplished, so go for it! Let us know if you get stuck with anything PennyLane related in the process :slight_smile: