How to incorporate a "black box" layer into a model? (Quantum Computing)

I’m trying to create a neural network where given the pure initial state of a quantum circuit (2D-vector), it spits out 2 numbers that would be essentially fed into a quantum computer to get results out. Currently, I have the step where I send a query to the quantum computer as a layer in the model. The quantum computer then will (in this case) spit out two numbers which I’d compare to the theoretical values.

As I was trying to implement this using the functional API, my custom layer was yelling at my because it did not like me using tf.unstack when it tried to build it by passing through a variable tensor. While this would probably be somewhat simple to fix, I ‘m concerned about how Qiskit would react to the values from the variable tensor or if it would work at all. Are there any workarounds or already implemented functions to achieve this?

Code:

Data Generation:

# TODO extend to complex amplitudes def generate_Hadamard_data(num_data): data = [] for _ in range(num_data): data_elem = preprocessing.normalize(np.random.rand(1, 2)).tolist()[0] data.append(data_elem) inv_sq2 = 1/np.sqrt(2) hadamard = np.array([[inv_sq2, inv_sq2], [inv_sq2, -1*inv_sq2]]) output = [np.matmul(hadamard, state).tolist() for state in data] data_ten = tf.constant(np.array(data).T, dtype=tf.float32) output_ten = tf.constant(np.array(output).T, dtype=tf.float32) print(output_ten) return tf.data.Dataset.from_tensor_slices(data_ten), tf.data.Dataset.from_tensor_slices(output_ten) 

Custom Layer:

class HadamardCircuitLayer(tf.keras.layers.Layer): ''' Takes in the initial data (state) and output (pulse params) from the neural network and runs it on the backend ibmq_armonk ''' def __init__(self, initial_data, shots=1024): super(HadamardCircuitLayer, self).__init__() # Unpack data data_iterator = initial_data.as_numpy_iterator() self.zero_coeffs = data_iterator.next() self.one_coeffs = data_iterator.next() self.shots = shots def build(self, shape): pass def call(self, output): custom_gate = Gate('custom_gate', 1, []) provider = IBMQ.get_provider(hub='ibm-q', group='open', project='main') backend = provider.get_backend('ibmq_armonk') post_q_list = [] for a, b, pulse_params in zip(self.zero_coeffs, self.one_coeffs, tf.unstack(output)): norm = np.sqrt(a**2 + b**2) qc = QuantumCircuit(1, 1) qc.initialize([a/norm, b/norm], 0) qc.append(custom_gate, [0]) qc.measure(0, 0) pul = pulse_params.numpy() with pulse.build(backend, name='custom') as my_schedule: pulse.play(Gaussian(duration=64, amp=pul[0], sigma=np.e**pul[1]), pulse.drive_channel(0)) qc.add_calibration(custom_gate, [0], my_schedule) qc = transpile(qc, backend) job = execute(qc, backend=backend, shots=self.shots) counts = job.result().get_counts() post_q_list.append(np.array([counts["0"]/self.shots, counts["1"]/self.shots])) output_qten = tf.constant(post_q_list, dtype=tf.float32) print(output_qten) return output_qten 

Implementation:

# I haven't gotten far since running into the unstack error x_train, y_train = generate_Hadamard_data(1) dataset = tf.data.Dataset.zip((x_train, y_train)) inputs = tf.keras.Input(shape=(2,)) x = tf.keras.layers.Dense(128, activation="linear")(inputs) x = tf.keras.layers.Dropout(0.2)(x) x = tf.keras.layers.Dense(128, activation="linear")(x) x = tf.keras.layers.Dense(2, activation="linear")(x) outputs = HadamardCircuitLayer(x_train)(x) model = tf.keras.Model(inputs=inputs, outputs=outputs, name="microwave_pulse_model")