Source code for pennylane.transforms.broadcast_expand
# Copyright 2018-2021 Xanadu Quantum Technologies Inc.
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# http://www.apache.org/licenses/LICENSE-2.0
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""This module contains the tape expansion function for expanding a
broadcasted tape into multiple tapes."""
import pennylane as qml
from pennylane.measurements import MidMeasureMP, SampleMP
from pennylane.tape import QuantumScript, QuantumScriptBatch
from pennylane.typing import PostprocessingFn
from .core import transform
def _split_operations(ops, num_tapes):
"""
Given a list of operators, return a list containing lists
of new operators with length num_tapes, with the parameters split.
"""
# for some reason pylint thinks "qml.ops" is a set
# pylint: disable=no-member
new_ops = [[] for _ in range(num_tapes)]
for op in ops:
# determine if any parameters of the operator are batched
if op.batch_size:
for b in range(num_tapes):
new_params = tuple(
p if qml.math.ndim(p) == op.ndim_params[j] else p[b]
for j, p in enumerate(op.data)
)
new_op = qml.ops.functions.bind_new_parameters(op, new_params)
new_ops[b].append(new_op)
else:
# no batching in the operator; don't copy
for b in range(num_tapes):
new_ops[b].append(op)
return new_ops
def null_postprocessing(results):
"""A postprocesing function returned by a transform that only converts the batch of results
into a result for a single ``QuantumTape``.
"""
return results[0]
[docs]@transform
def broadcast_expand(tape: QuantumScript) -> tuple[QuantumScriptBatch, PostprocessingFn]:
r"""Expand a broadcasted tape into multiple tapes
and a function that stacks and squeezes the results.
.. warning::
Currently, not all templates have been updated to support broadcasting.
Args:
tape (QNode or QuantumTape or Callable): Broadcasted tape to be expanded
Returns:
qnode (QNode) or quantum function (Callable) or tuple[List[QuantumTape], function]:
The transformed circuit as described in :func:`qml.transform <pennylane.transform>`.
- If the input is a QNode, the broadcasted input QNode
that computes the QNode output serially with multiple circuit evaluations and
stacks (and squeezes) the results into one batch of results.
- If the input is a tape, a tuple containing a list of generated tapes, together with
a post-processing function. The number of tapes matches the broadcasting dimension
of the input tape, and the results from the evaluated tapes are stacked and squeezed
together in the post-processing function.
This expansion function is used internally whenever a device does not
support broadcasting.
**Example**
We may use ``broadcast_expand`` on a ``QNode`` to separate it
into multiple calculations. For this we will provide ``qml.RX`` with
the ``ndim_params`` attribute that allows the operation to detect
broadcasting, and set up a simple ``QNode`` with a single operation and
returned expectation value:
>>> from pennylane import numpy as np
>>> qml.RX.ndim_params = (0,)
>>> dev = qml.device("default.qubit", wires=1)
>>> @qml.qnode(dev)
>>> def circuit(x):
... qml.RX(x, wires=0)
... return qml.expval(qml.Z(0))
We can then call ``broadcast_expand`` on the QNode and store the
expanded ``QNode``:
>>> expanded_circuit = qml.transforms.broadcast_expand(circuit)
Let's use the expanded QNode and draw it for broadcasted parameters
with broadcasting axis of length ``3`` passed to ``qml.RX``:
>>> x = np.array([0.2, 0.6, 1.0], requires_grad=True)
>>> print(qml.draw(expanded_circuit)(x))
0: ──RX(0.20)─┤ <Z>
0: ──RX(0.60)─┤ <Z>
0: ──RX(1.00)─┤ <Z>
Executing the expanded ``QNode`` results in three values, corresponding
to the three parameters in the broadcasted input ``x``:
>>> expanded_circuit(x)
tensor([0.98006658, 0.82533561, 0.54030231], requires_grad=True)
We also can call the transform manually on a tape:
>>> ops = [qml.RX(np.array([0.2, 0.6, 1.0], requires_grad=True), wires=0)]
>>> measurements = [qml.expval(qml.Z(0))]
>>> tape = qml.tape.QuantumTape(ops, measurements)
>>> tapes, fn = qml.transforms.broadcast_expand(tape)
>>> tapes
[<QuantumTape: wires=[0], params=1>, <QuantumTape: wires=[0], params=1>, <QuantumTape: wires=[0], params=1>]
>>> fn(qml.execute(tapes, qml.device("default.qubit", wires=1), None))
tensor([0.98006658, 0.82533561, 0.54030231], requires_grad=True)
"""
if tape.batch_size is None:
return (tape,), null_postprocessing
has_postselect = any(
op.postselect is not None for op in tape.operations if isinstance(op, MidMeasureMP)
)
has_sample = any(isinstance(op, SampleMP) for op in tape.measurements)
if has_postselect and has_sample:
raise ValueError(
"Returning qml.sample is not supported when using post-selected mid-circuit measurements and parameters broadcasting."
)
num_tapes = tape.batch_size
new_ops = _split_operations(tape.operations, num_tapes)
output_tapes = tuple(
qml.tape.QuantumScript(
ops, tape.measurements, shots=tape.shots, trainable_params=tape.trainable_params
)
for ops in new_ops
)
def processing_fn(results: qml.typing.ResultBatch) -> qml.typing.Result:
# closure variables: tape.shots, tape.batch_size, tape.measurements
# The shape of the results should be as follows: results[s][m][b], where s is the shot
# vector index, m is the measurement index, and b is the batch index. The shape that
# the processing function receives is results[b][s][m].
if tape.shots.has_partitioned_shots:
if len(tape.measurements) > 1:
return tuple(
tuple(
qml.math.stack([results[b][s][m] for b in range(tape.batch_size)])
for m in range(len(tape.measurements))
)
for s in range(tape.shots.num_copies)
)
# Only need to transpose results[b][s] -> results[s][b]
return tuple(
qml.math.stack([results[b][s] for b in range(tape.batch_size)])
for s in range(tape.shots.num_copies)
)
if len(tape.measurements) > 1:
# Only need to transpose results[b][m] -> results[m][b]
return tuple(
qml.math.stack([results[b][m] for b in range(tape.batch_size)])
for m in range(len(tape.measurements))
)
return qml.math.stack(results)
return output_tapes, processing_fn
_modules/pennylane/transforms/broadcast_expand
Download Python script
Download Notebook
View on GitHub