Apply a quantum circuit on a specific dimension of a tensor

Hi,

I was wondering how to apply a quantum circuit for instance on the last dimension of a 3D or 4D tensor. For instance, let’s say I have an input tensor of size (32,32,3) and want to apply a quantum circuit as a KerasLayer (with tensorflow but maybe interesting to also mention the TorchLayer) to get an output of size (32,32,1).

Thanks for your support.

Hi @cnada,

If I understand your question correctly, you have a QNode that takes 3-dimensional inputs and returns 1-dimensional outputs, and you want to apply this over a tensor of rank greater than 2?

This can be done with reshaping:
(TF version)

import pennylane as qml
import tensorflow as tf

input_dim = 3
output_dim = 1
other_dim = 32

n_layers = 2

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

@qml.qnode(dev)
def qnode(inputs, weights):
    qml.templates.AngleEmbedding(inputs, wires=range(input_dim))
    qml.templates.StronglyEntanglingLayers(weights, wires=range(input_dim))
    return [qml.expval(qml.PauliZ(i)) for i in range(output_dim)]
    
weight_shapes = {"weights": (n_layers, input_dim, 3)}

qlayer = qml.qnn.KerasLayer(qnode, weight_shapes, output_dim=output_dim)

x = tf.ones((other_dim, other_dim, input_dim))
x = tf.reshape(x, (other_dim * other_dim, input_dim))
tf.reshape(qlayer(x), (other_dim, other_dim, output_dim)).shape

You can also use the Reshape layer in Keras.

(torch version)

import pennylane as qml
import torch

input_dim = 3
output_dim = 1
other_dim = 32

n_layers = 2

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

@qml.qnode(dev)
def qnode(inputs, weights):
    qml.templates.AngleEmbedding(inputs, wires=range(input_dim))
    qml.templates.StronglyEntanglingLayers(weights, wires=range(input_dim))
    return [qml.expval(qml.PauliZ(i)) for i in range(output_dim)]
    
weight_shapes = {"weights": (n_layers, input_dim, 3)}

qlayer = qml.qnn.TorchLayer(qnode, weight_shapes)

x = torch.ones((other_dim, other_dim, input_dim))
x = torch.reshape(x, (other_dim * other_dim, input_dim))
torch.reshape(qlayer(x), (other_dim, other_dim, output_dim)).shape

In general, if your QNode expects inputs to have a given shape, then the tensor fed into the corresponding KerasLayer or TorchLayer should have shape (batch_size, *shape).

@Tom_Bromley yes. Well actually I meant more say we have an input tensor of shape (None, 32, 32, 3) (think of as a batch of images), and the output will be (None, 32, 32, 1). So basically the qlayer will apply only in the last dimension of the tensor. But from your answer it seems reshaping to (-1,3) is the only way. I thought maybe there was another way.