It appears that the qml.qnn.TorchLayer()
although allow multiple weights, it only allows the qnode has one inputs
argument, and requires packing multiple inputs input a torch array and using index to access them. In the following example, I would like to have the circuit()
qnode accept two arrays of different length as inputs, so that the two QAOAEmbedding
can accept different input array length, is it possible to do that?
import pennylane as qml
import torch
import torch.nn as nn
n_qbits = 8
m = 4
layer = 2
qdev = qml.device("default.qubit", n_qbits)
@qml.qnode(qdev, interface="torch")
def circuit(inputs, weights_a, weights_b):
qml.QAOAEmbedding(inputs[0], weights=weights_a, wires=range(n_qbits))
qml.adjoint(qml.QAOAEmbedding(inputs[1], weights=weights_b, wires=range(n_qbits)))
return qml.probs(wires=range(m))
class QSim(nn.Module):
def __init__(self):
super().__init__()
self.qlayer = qml.qnn.TorchLayer(
circuit,
weight_shapes={
"weights_a": (layer, 2 * n_qbits),
"weights_b": (layer, 2 * n_qbits),
},
)
def forward(self, a, b):
return self.qlayer(torch.stack([a, b], dim=1))[:, 0]
model = QSim().to(device)
# batch num = 8, array size = 4
model(torch.randn(8, 4).to(device), torch.randn(8, 4).to(device))
# tensor([0.0367, 0.0399, 0.0973, 0.0499, 0.1242, 0.0184, 0.1547, 0.0551],
# device='cuda:0', grad_fn=<SelectBackward0>)