Issue when optimizing a quimb quantum circuit through tensorflow.

[LucasAugustusvd]

What is your issue?

I am trying to optimize a quantum circuit I made with quimb using tensorflow but I run into issues with how some operations in quimb are not compatible with tensorlfows graph method. For instance I incorporate CNOT gates in my circuit but tensorflow seems to want to use tf.max instead of the normal numpy operation that quimb uses. Does anyone see any glaring issue with my code that I am missing?

Model:

import tensorflow as tf
import numpy as np
import quimb.tensor as qtn

class MPSQuantumLayer(tf.keras.layers.Layer):
    def __init__(self, n_qubits, n_layers, max_bond=10, cutoff=1e-5):
        super(MPSQuantumLayer, self).__init__()
        self.n_qubits = n_qubits
        self.n_layers = n_layers
        self.max_bond = max_bond
        self.cutoff = cutoff
    
    def build(self, input_shape):
        self.params = self.add_weight(shape=(self.n_layers + 1, self.n_qubits, 3), 
                                      initializer='random_normal', trainable=True)
        self.inputs = self.add_weight(shape=(self.n_layers, self.n_qubits), 
                                      initializer='random_normal', trainable=True)
    @tf.function
    def call(self, inputs):
        circuit = qtn.CircuitMPS(self.n_qubits, max_bond=self.max_bond, cutoff=self.cutoff)
        for j in range(self.n_layers):
            for i in range(self.n_qubits):
                circuit.apply_gate('RX', self.params[j, i, 0], i, gate_round=j)
                circuit.apply_gate('RY', self.params[j, i, 1], i, gate_round=j)
                circuit.apply_gate('RZ', self.params[j, i, 2], i, gate_round=j)
                if i < self.n_qubits - 1:
                    circuit.apply_gate('CNOT', i, i + 1, gate_round=j)
                circuit.apply_gate('RX', self.inputs[j, i], i, gate_round=j)

        for i in range(self.n_qubits):
            circuit.apply_gate('RX', self.params[self.n_layers, i, 0], i)
            circuit.apply_gate('RY', self.params[self.n_layers, i, 1], i)
            circuit.apply_gate('RZ', self.params[self.n_layers, i, 2], i)

        # Calculate expectation values for Pauli X and Z
        G_X = np.array([[0, 1], [1, 0]], dtype='complex128')
        G_Z = np.array([[1, 0], [0, -1]], dtype='complex128')
        expectations_X = [circuit.local_expectation(G_X, where=i, optimize='auto-hq').real for i in range(self.n_qubits)]
        expectations_Z = [circuit.local_expectation(G_Z, where=i, optimize='auto-hq').real for i in range(self.n_qubits)]
        
        return tf.cast([tf.cast([expectations_X[i], expectations_Z[i]], axis=-1) for i in range(self.n_qubits)], axis=1)

    def compute_output_shape(self, input_shape):
        # Output shape includes real parts of expectations for X and Z for each qubit
        return (input_shape[0], self.n_qubits, 2)


        

class AlternatingLayer(tf.keras.layers.Layer):
    def __init__(self, n_actions):
        super(AlternatingLayer, self).__init__()
        self.n_actions = n_actions
        # Initialize weights with shape (n_qubits, 2) to match input shape
        self.w = self.add_weight(
            shape=(n_actions, 2),  # Updated shape to match the input's second and third dimensions
            initializer='random_normal',
            trainable=True
        )

    def call(self, inputs):
        # Reshape inputs to (batch_size, n_qubits * 2)
        inputs = tf.reshape(inputs, (tf.shape(inputs)[0], -1))
        # Reshape weights to (n_qubits * 2, n_actions)
        w_reshaped = tf.reshape(self.w, (-1, self.n_actions))

        # Perform matrix multiplication
        output = tf.matmul(inputs, w_reshaped)

        return output



def generate_model_policy(n_qubits, n_layers, n_actions):
    """Generates a Keras model using MPS and an Alternating layer."""
    input_tensor = tf.keras.Input(shape=(n_qubits, ), dtype=tf.dtypes.float32, name='input')  # Shape adjusted for 2 outputs per qubit
    quantum_layer = MPSQuantumLayer(n_qubits, n_layers)(input_tensor)
    alternating_layer = AlternatingLayer(n_actions)(quantum_layer)
    output = tf.keras.layers.Dense(1, activation='sigmoid')(alternating_layer)
    model = tf.keras.Model(inputs=input_tensor, outputs=output)

    return model




def make_discriminator_model_conv(chopsize):
    model = tf.keras.Sequential()
    model.add(tf.keras.layers.Convolution1D(64,10,padding='same',input_shape=(chopsize,1)))
    model.add(tf.keras.layers.LeakyReLU(0.2))
    model.add(tf.keras.layers.Convolution1D(128,10,padding='same'))
    model.add(tf.keras.layers.LeakyReLU(0.2))
    model.add(tf.keras.layers.Convolution1D(128,10,padding='same'))
    model.add(tf.keras.layers.LeakyReLU(0.2))
    model.add(tf.keras.layers.Flatten())
    model.add(tf.keras.layers.Dense(32))
    model.add(tf.keras.layers.LeakyReLU(0.2))
    model.add(tf.keras.layers.Dropout(0.5))
    model.add(tf.keras.layers.Dense(1))
    return model
    
    #Code snippet for how I train this model:
    
    @tf.function
    def train_generator(self, gan_instance:tf.keras.Model, critic:tf.keras.Model, noise: tf.Tensor):
        """ Update the generator parameter with its optimizer"""

        with tf.GradientTape() as tape:
            generated_images = gan_instance(noise, training=False)
            generated_images = tf.expand_dims(generated_images, -1)
            fake_output = critic(generated_images, training=True)

            gen_loss = self.wasserstein_loss_generator(fake_output)

        gradients_of_generator = tape.gradient(
            gen_loss, gan_instance.trainable_variables
        )
        self.generator_optimizer.apply_gradients(
            zip(gradients_of_generator, gan_instance.trainable_variables)
        )

    @tf.function
    def train_critic(self, gan_instance:tf.keras.Model, critic:tf.keras.Model, noise: tf.Tensor, images: tf.Tensor):
        """Update the generator parameter with its optimizer

        :param noise: Input of the generator
        :param images: Real images as input for the discriminator.
        """

        with tf.GradientTape() as tape:
            generated_images = gan_instance(noise, training=False)
            generated_images = tf.expand_dims(generated_images, -1)
            real_output = critic(images, training=True)
            fake_output = critic(generated_images, training=True)
            disc_loss = self.wasserstein_loss_critic(real_output, fake_output)
            penalty_loss = self.gradient_penalty(critic, generated_images, images)
            disc_loss += penalty_loss * self.gradient_penalty_weight

        gradients_of_critic = tape.gradient(disc_loss, critic.trainable_variables)
        self.discriminator_optimizer.apply_gradients(
            zip(gradients_of_critic, critic.trainable_variables)
        )

    def train(self):
        print('Run: ', self.run_id, ' has started')
        count = 0
        lowest_wass = 1e5
        lowest_acf_abs = 1e5
        lowest_acf_nonabs = 1e5
        lowest_leverage = 1e5
        noise_cst = tf.random.uniform([self.train_time_series.shape[0], self.noise_dim], minval=0, maxval=2*np.pi, seed = 0)
        train_dataset = tf.data.Dataset.from_tensor_slices(self.train_time_series).shuffle(self.BUFFER_SIZE).batch(self.BATCH_SIZE * self.nb_steps_update_critic)
        lowest_wass_epoch, lowest_acf_abs_epoch, lowest_acf_nonabs_epoch, lowest_leverage_epoch = 0,0,0,0
        for epoch in tqdm(range(self.epochs)):
            for image_batch in train_dataset:
                batches_discriminator = tf.data.Dataset.from_tensor_slices(image_batch).batch(self.BATCH_SIZE)
                for discriminator_batch in batches_discriminator:
                    noise = tf.random.uniform([discriminator_batch.shape[0], self.noise_dim], minval=0, maxval=2*np.pi)

                    discriminator_batch = tf.expand_dims(discriminator_batch, -1)

                    self.train_critic(self.generator, self.discriminator, noise, discriminator_batch)

                noise = tf.random.uniform([self.BATCH_SIZE, self.noise_dim],minval=0, maxval=2*np.pi)
                self.train_generator(self.generator, self.discriminator, noise)