qml.batch_params

batch_params(tape, all_operations=False)

Transform a QNode to support an initial batch dimension for operation parameters.

Note

This transform will create multiple circuits inside the QNode, one per batch dimension. As a result, it is both simulator and hardware compatible. When using a simulator device, however, this means that a separate simulation will be performed per batch dimension.

Warning

Currently, not all templates have been updated to support a batch dimension. If you run into an error attempting to use a template with this transform, please open a GitHub issue detailing the error.

Parameters

• qnode (pennylane.QNode or QuantumTape) – quantum tape or QNode to add a batch dimension to

• all_operations (bool) – If True, a batch dimension will be added to all operations in the QNode, rather than just trainable QNode parameters.

Returns

Function which accepts the same arguments as the QNode, however the first dimension of each argument will be treated as a batch dimension. The function output will also contain an initial batch dimension.

Return type

func

Example

Consider the following circuit:

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

@qml.batch_params
@qml.qnode(dev)
def circuit(x, weights):
    qml.RX(x, wires=0)
    qml.RY(0.2, wires=1)
    qml.templates.StronglyEntanglingLayers(weights, wires=[0, 1, 2])
    return qml.expval(qml.Hadamard(0))

The qml.batch_params decorator allows us to pass arguments x and weights that have a batch dimension. For example,

>>> batch_size = 3
>>> x = np.linspace(0.1, 0.5, batch_size)
>>> rng = np.random.default_rng(seed=1234)
>>> weights = rng.random((batch_size, 10, 3, 3), requires_grad=True)

If we evaluate the QNode with these inputs, we will get an output of shape (batch_size,):

>>> circuit(x, weights)
tensor([ 0.00800498,  0.2735391 , -0.24395442], requires_grad=True)

QNodes with a batch dimension remain fully differentiable:

>>> cost_fn = lambda x, weights: np.sum(circuit(x, weights))
>>> cost_fn(x, weights)
tensor(0.03758966, requires_grad=True)
>>> qml.grad(cost_fn)(x, weights)[0]
array([-0.30262974,  0.06320878,  0.00811555])

If we pass the all_operations argument, we can specify that all operation parameters in the transformed QNode, regardless of whether they are QNode input parameters, have a batch dimension:

@qml.batch_params(all_operations=True)
@qml.qnode(dev)
def circuit(x, weights):
    qml.RX(x, wires=0)
    qml.RY([0.2, 0.2, 0.2], wires=1)
    qml.templates.StronglyEntanglingLayers(weights, wires=[0, 1, 2])
    return qml.expval(qml.Hadamard(0))

>>> cost_fn = lambda x, weights: np.sum(circuit(x, weights))
>>> weights.requires_grad = False
>>> cost_fn(x, weights)
tensor(0.03758966, requires_grad=True)
>>> qml.grad(cost_fn)(x, weights)[0]
-0.30262974103192636

code/api/pennylane.batch_params

 

Download Python script

 

Download Notebook

 

View on GitHub