Optimize with qml.state() and autograd

Hello, I’m trying to calculate inner products of quantum states directly from qml.state() and build kernel matrix. My optimizer is qml.AdamOptimizer().

And I got TypeError: must be real number, not ArrayBox when assign quantum state vectors to an numpy array element wise. It seems like the returned state is ArrayBox type and cannot be assigned to a numpy array element.

Is there any way to solve this?:grinning:
I also tried to use ArrayBox._value to assign. However in this way the weights just stayed the same and didn’t update.

Below are my codes:

import pennylane as qml
from pennylane import numpy as np

n_qubits = 5
n_layers = 2
dev = qml.device("default.qubit", wires=n_qubits)

from pennylane.templates import QAOAEmbedding
@qml.qnode(dev, diff_method='backprop')
def circuit(weights, feature):
    QAOAEmbedding(features=feature, weights=weights, wires=range(n_qubits))
    return qml.state()

from pennylane.init import qaoa_embedding_uniform
weights = qaoa_embedding_uniform(n_layers,n_qubits)
#print(weights)

N = 10
d = 4
X = np.random.uniform(size=(N,d))
#print(X)

# map x to quantum state |x>
def encoder(weights, X):
    M = X.shape[0]
    features = np.empty((M,2**n_qubits),dtype='complex_')
    for i in range(M):
        # failed here
        features[i] = circuit(weights,X[i])
        # features[i] = circuit(weights,X[i])._value
    return features
# kernel matrix
def matrix(features):
    # <x_i|x_j>
    t1 = features @ features.conj().T
    # |<x_i|x_j>|^2
    t2 = np.multiply(t1,t1.conj())
    return np.absolute(t2)

def cost_naive(weights,X=None):
    features = encoder(weights,X)
    K = matrix(features)
    return K.sum()

optim = qml.AdamOptimizer()
for i in range(10):
    weights, loss = optim.step_and_cost(cost_naive,weights,X=X)
    print(weights[0])
    print(loss)

Hi @Zedd!

The issue here is that Autograd does not support index assignment; in particular the line

features[i] = circuit(weights,X[i])

where the output of circuit is being assigned to the ith element of features.

This is a common ‘gotcha’ of Autograd; the Autograd tutorial lists other parts of NumPy/SciPy that are not supported by Autograd.

But that’s all good! Rather than using array index assignment, we can rewrite this to use lists instead:

features = []
for i in range(M):
    features.append(circuit(weights, X[i]))
return np.stack(features)

Here, we are creating a list containing the circuit results, and then ‘stacking’ them via np.stack to convert it into an array.

There is one more change needed to get it working — we need to replace array.conj() with np.conj(array) (this former is another NumPy feature seemingly not supported by Autograd). With these two changes, the script now works:

import pennylane as qml
from pennylane import numpy as np
from pennylane.init import qaoa_embedding_uniform
from pennylane.templates import QAOAEmbedding

n_qubits = 5
n_layers = 2

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

@qml.qnode(dev, diff_method="backprop")
def circuit(weights, feature):
    QAOAEmbedding(features=feature, weights=weights, wires=range(n_qubits))
    return qml.state()

weights = qaoa_embedding_uniform(n_layers, n_qubits)

N = 10
d = 4
X = np.random.uniform(size=(N, d))

# map x to quantum state |x>
def encoder(weights, X):
    M = X.shape[0]
    features = []
    for i in range(M):
        features.append(circuit(weights, X[i]))
    return np.stack(features)

# kernel matrix
def matrix(features):
    # <x_i|x_j>
    t1 = features @ np.conj(features).T
    # |<x_i|x_j>|^2
    t2 = np.multiply(t1, np.conj(t1))
    return np.absolute(t2)

def cost_naive(weights, X=None):
    features = encoder(weights, X)
    K = matrix(features)
    return K.sum()

optim = qml.AdamOptimizer()
for i in range(10):
    weights, loss = optim.step_and_cost(cost_naive, weights, X=X)
    print(weights[0])

Let me know if that works for you, or if you have any other questions!

Wow, thank you @josh, it works. And thanks for explaining it in every detail.

1 Like

No worries @Zedd :slight_smile: