Train a Parametrized Quantum Circuit with a loss function using the samples measurement

General Discussion
,
#1

Hello,
While trying to implement this article Implementable Quantum Classifier for Nonlinear Data i syumble upon a problem.
I cannot find a way to train my circuit using a loss function that is defined with the amplitudes of the final state and not on average measurement over every Qubits.

Here are the relevant parts of my code :

def gaussian_kernel(x,y,sigma) :
return np.exp(-np.abs(x-y)**2/2*sigma)

def loss_func(output,k):
l_outuput = list(output)
dico_res = {}
count = 0
for el in output[0] :
I_reg = el.numpy()[:N_qubit_F]
key = sum([2**i for i in I_reg])
if key not in dico_res.keys():
dico_res[key]=1
else:
dico_res[key]+=1
count += 1
loss = 0
for key in dico_res.keys():
loss += gaussian_kernel(dico_res[key],dico_res[key],sigma)
loss += gaussian_kernel(dico_res[key],0,sigma)
if k in dico_res.keys():
loss -= gaussian_kernel(dico_res[key],0,sigma)
loss += gaussian_kernel(dico_res[key],1,sigma)
loss += 1
return loss/count

optimizer = tf.keras.optimizers.Adam(1e-4)

model = tf.keras.Sequential([
tf.keras.layers.Input(shape=(), dtype=tf.string),
tfq.layers.PQC(circuit,[cirq.Z(cirq.GridQubit(0,0))])
])

@tf.function
def train_step(circuit_input):
with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape:
sample_layer = tfq.layers.Sample()
output = sample_layer(circuit_input[0] + circuit, repetitions=100)

    loss = loss_func(output,k)
print(loss)

gradients = gen_tape.gradient(
    loss, symbols)
print(gradients)

optimizer.apply_gradients(
    zip(gradients, model.trainable_variables))

return loss

train_step(tfq.convert_to_tensor([circuit_input]))

The circuit in itself is a paramtetrized one.
Thanks a lot for your help.

#2

/cc @Michael_Broughton