Source code for pennylane.gradients.parameter_shift_cv
# 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 functions for computing the parameter-shift gradient
of a CV-based quantum tape.
"""
import itertools
import warnings
from functools import partial
import numpy as np
import pennylane as qml
from pennylane import transform
from pennylane.gradients.gradient_transform import (
_contract_qjac_with_cjac,
_validate_gradient_methods,
choose_trainable_params,
)
from pennylane.measurements import (
ExpectationMP,
MeasurementProcess,
ProbabilityMP,
StateMP,
VarianceMP,
)
from pennylane.tape import QuantumScript, QuantumScriptBatch
from pennylane.transforms.tape_expand import expand_invalid_trainable
from pennylane.typing import PostprocessingFn
from .finite_difference import finite_diff
from .general_shift_rules import generate_shifted_tapes, process_shifts
from .gradient_transform import _no_trainable_grad
from .parameter_shift import _get_operation_recipe, expval_param_shift
# pylint: disable=protected-access,too-many-arguments,too-many-statements,too-many-branches,unused-argument
def _grad_method_cv(tape, idx):
"""Determine the best CV parameter-shift gradient recipe for a given
parameter index of a tape.
Args:
tape (.QuantumTape): input tape
idx (int): positive integer corresponding to the parameter location
on the tape to inspect
Returns:
str: a string containing either ``"A"`` (for first-order analytic method),
``"A2"`` (second-order analytic method), ``"F"`` (finite differences),
or ``"0"`` (constant parameter).
"""
par_info = tape.par_info[idx]
op = par_info["op"]
if op.grad_method in (None, "F"):
return op.grad_method
if op.grad_method != "A":
raise ValueError(f"Operation {op} has unknown gradient method {op.grad_method}")
# Operation supports the CV parameter-shift rule.
# Create an empty list to store the 'best' partial derivative method
# for each observable
best = []
for m in tape.measurements:
if isinstance(m, ProbabilityMP) or (m.obs.ev_order not in (1, 2)):
# Higher-order observables (including probability) only support finite differences.
best.append("F")
continue
# get the set of operations betweens the operation and the observable
op_or_mp = tape[par_info["op_idx"]]
ops_between = tape.graph.nodes_between(op_or_mp, m)
if not ops_between:
# if there is no path between the operation and the observable,
# the operator has a zero gradient.
best.append("0")
continue
# For parameter-shift compatible CV gates, we need to check both the
# intervening gates, and the type of the observable.
best_method = "A"
ops_between = [o.obs if isinstance(o, MeasurementProcess) else o for o in ops_between]
if any(not k.supports_heisenberg for k in ops_between):
# non-Gaussian operators present in-between the operation
# and the observable. Must fallback to numeric differentiation.
best_method = "F"
elif m.obs.ev_order == 2:
if isinstance(m, ExpectationMP):
# If the observable is second-order, we must use the second-order
# CV parameter shift rule
best_method = "A2"
elif isinstance(m, VarianceMP):
# we only support analytic variance gradients for
# first-order observables
best_method = "F"
best.append(best_method)
if all(k == "0" for k in best):
# if the operation is independent of *all* observables
# in the circuit, the gradient will be 0
return "0"
if "F" in best:
# one non-analytic observable path makes the whole operation
# gradient method fallback to finite-difference
return "F"
if "A2" in best:
# one second-order observable makes the whole operation gradient
# require the second-order parameter-shift rule
return "A2"
return "A"
def _find_gradient_methods_cv(tape, trainable_param_indices):
"""Find the best gradient methods for each parameter."""
return {
idx: _grad_method_cv(tape, tape.trainable_params[idx]) for idx in trainable_param_indices
}
def _gradient_analysis_and_validation_cv(tape, method, trainable_param_indices):
"""Find the best gradient methods for each parameter. Subsequently, validate
the gradient methods and return diff_methods."""
diff_methods = _find_gradient_methods_cv(tape, trainable_param_indices)
_validate_gradient_methods(tape, method, diff_methods)
return diff_methods
def _transform_observable(obs, Z, device_wires):
"""Apply a Gaussian linear transformation to an observable.
Args:
obs (.Observable): observable to transform
Z (array[float]): Heisenberg picture representation of the linear transformation
device_wires (.Wires): wires on the device the transformed observable is to be
measured on
Returns:
.Observable: the transformed observable
"""
# Get the Heisenberg representation of the observable
# in the position/momentum basis. The returned matrix/vector
# will have been expanded to act on the entire device.
if obs.ev_order > 2:
raise NotImplementedError("Transforming observables of order > 2 not implemented.")
A = obs.heisenberg_obs(device_wires)
if A.ndim != obs.ev_order:
raise ValueError(
"Mismatch between the polynomial order of observable and its Heisenberg representation"
)
# transform the observable by the linear transformation Z
A = A @ Z
if A.ndim == 2:
A = A + A.T
# TODO: if the A matrix corresponds to a known observable in PennyLane,
# for example qml.QuadX, qml.QuadP, qml.NumberOperator, we should return that
# instead. This will allow for greater device compatibility.
return qml.PolyXP(A, wires=device_wires)
def var_param_shift(tape, dev_wires, argnum=None, shifts=None, gradient_recipes=None, f0=None):
r"""Partial derivative using the first-order or second-order parameter-shift rule of a tape
consisting of a mixture of expectation values and variances of observables.
Expectation values may be of first- or second-order observables,
but variances can only be taken of first-order variables.
.. warning::
This method can only be executed on devices that support the
:class:`~.PolyXP` observable.
Args:
tape (.QuantumTape): quantum tape to differentiate
dev_wires (.Wires): wires on the device the parameter-shift method is computed on
argnum (int or list[int] or None): Trainable parameter indices to differentiate
with respect to. If not provided, the derivative with respect to all
trainable indices are returned.
shifts (list[tuple[int or float]]): List containing tuples of shift values.
If provided, one tuple of shifts should be given per trainable parameter
and the tuple should match the number of frequencies for that parameter.
If unspecified, equidistant shifts are assumed.
gradient_recipes (tuple(list[list[float]] or None)): List of gradient recipes
for the parameter-shift method. One gradient recipe must be provided
per trainable parameter.
f0 (tensor_like[float] or None): Output of the evaluated input tape. If provided,
and the gradient recipe contains an unshifted term, this value is used,
saving a quantum evaluation.
Returns:
tuple[list[QuantumTape], function]: A tuple containing a
list of generated tapes, together with a post-processing
function to be applied to the results of the evaluated tapes
in order to obtain the Jacobian matrix.
"""
argnum = argnum or tape.trainable_params
# Determine the locations of any variance measurements in the measurement queue.
var_mask = [isinstance(m, VarianceMP) for m in tape.measurements]
var_idx = np.where(var_mask)[0]
# Get <A>, the expectation value of the tape with unshifted parameters.
expval_tape = tape.copy(copy_operations=True)
# Convert all variance measurements on the tape into expectation values
for i in var_idx:
obs = expval_tape.measurements[i].obs
expval_tape._measurements[i] = qml.expval(op=obs)
gradient_tapes = [expval_tape]
# evaluate the analytic derivative of <A>
pdA_tapes, pdA_fn = expval_param_shift(expval_tape, argnum, shifts, gradient_recipes, f0)
gradient_tapes.extend(pdA_tapes)
# Store the number of first derivative tapes, so that we know
# the number of results to post-process later.
tape_boundary = len(pdA_tapes) + 1
expval_sq_tape = tape.copy(copy_operations=True)
for i in var_idx:
# We need to calculate d<A^2>/dp; to do so, we replace the
# observables A in the queue with A^2.
obs = expval_sq_tape.measurements[i].obs
# CV first-order observable
# get the heisenberg representation
# This will be a real 1D vector representing the
# first-order observable in the basis [I, x, p]
A = obs._heisenberg_rep(obs.parameters)
# take the outer product of the heisenberg representation
# with itself, to get a square symmetric matrix representing
# the square of the observable
obs = qml.PolyXP(np.outer(A, A), wires=obs.wires)
expval_sq_tape._measurements[i] = qml.expval(op=obs)
# Non-involutory observables are present; the partial derivative of <A^2>
# may be non-zero. Here, we calculate the analytic derivatives of the <A^2>
# observables.
pdA2_tapes, pdA2_fn = second_order_param_shift(
expval_sq_tape, dev_wires, argnum, shifts, gradient_recipes
)
gradient_tapes.extend(pdA2_tapes)
def processing_fn(results):
mask = qml.math.convert_like(qml.math.reshape(var_mask, [-1, 1]), results[0])
f0 = qml.math.expand_dims(results[0], -1)
pdA = pdA_fn(results[1:tape_boundary])
pdA2 = pdA2_fn(results[tape_boundary:])
# return d(var(A))/dp = d<A^2>/dp -2 * <A> * d<A>/dp for the variances (mask==True)
# d<A>/dp for plain expectations (mask==False)
pdA = qml.math.stack(pdA)
return qml.math.where(mask, pdA2 - 2 * f0 * pdA, pdA)
return gradient_tapes, processing_fn
def second_order_param_shift(tape, dev_wires, argnum=None, shifts=None, gradient_recipes=None):
r"""Generate the second-order CV parameter-shift tapes and postprocessing methods required
to compute the gradient of a gate parameter with respect to an
expectation value.
.. note::
The 2nd order method can handle also first-order observables, but
1st order method may be more efficient unless it's really easy to
experimentally measure arbitrary 2nd order observables.
.. warning::
The 2nd order method can only be executed on devices that support the
:class:`~.PolyXP` observable.
Args:
tape (.QuantumTape): quantum tape to differentiate
dev_wires (.Wires): wires on the device the parameter-shift method is computed on
argnum (int or list[int] or None): Trainable parameter indices to differentiate
with respect to. If not provided, the derivative with respect to all
trainable indices are returned.
shifts (list[tuple[int or float]]): List containing tuples of shift values.
If provided, one tuple of shifts should be given per trainable parameter
and the tuple should match the number of frequencies for that parameter.
If unspecified, equidistant shifts are assumed.
gradient_recipes (tuple(list[list[float]] or None)): List of gradient recipes
for the parameter-shift method. One gradient recipe must be provided
per trainable parameter.
Returns:
tuple[list[QuantumTape], function]: A tuple containing a
list of generated tapes, together with a post-processing
function to be applied to the results of the evaluated tapes
in order to obtain the Jacobian matrix.
"""
argnum = argnum or list(tape.trainable_params)
gradient_recipes = gradient_recipes or [None] * len(argnum)
gradient_tapes = []
shapes = []
obs_indices = []
gradient_values = []
for idx, _ in enumerate(tape.trainable_params):
t_idx = list(tape.trainable_params)[idx]
op = tape.par_info[t_idx]["op"]
if idx not in argnum:
# parameter has zero gradient
shapes.append(0)
obs_indices.append([])
gradient_values.append([])
continue
shapes.append(1)
# get the gradient recipe for the trainable parameter
arg_idx = argnum.index(idx)
recipe = gradient_recipes[arg_idx]
if recipe is not None:
recipe = process_shifts(np.array(recipe))
else:
op_shifts = None if shifts is None else shifts[arg_idx]
recipe = _get_operation_recipe(tape, idx, shifts=op_shifts)
coeffs, multipliers, op_shifts = recipe.T
if len(op_shifts) != 2:
# The 2nd order CV parameter-shift rule only accepts two-term shifts
raise NotImplementedError(
"Taking the analytic gradient for order-2 operators is "
f"unsupported for operation {op} which has a "
"gradient recipe of more than two terms."
)
shifted_tapes = generate_shifted_tapes(tape, idx, op_shifts, multipliers)
# evaluate transformed observables at the original parameter point
# first build the Heisenberg picture transformation matrix Z
Z0 = op.heisenberg_tr(dev_wires, inverse=True)
Z2 = shifted_tapes[0].par_info[t_idx]["op"].heisenberg_tr(dev_wires)
Z1 = shifted_tapes[1].par_info[t_idx]["op"].heisenberg_tr(dev_wires)
# derivative of the operation
Z = Z2 * coeffs[0] + Z1 * coeffs[1]
Z = Z @ Z0
# conjugate Z with all the descendant operations
B = np.eye(1 + 2 * len(dev_wires))
B_inv = B.copy()
succ = tape.graph.descendants_in_order((op,))
operation_descendents = itertools.filterfalse(
lambda obj: isinstance(obj, MeasurementProcess), succ
)
observable_descendents = filter(lambda obj: isinstance(obj, MeasurementProcess), succ)
for BB in operation_descendents:
if not BB.supports_heisenberg:
# if the descendant gate is non-Gaussian in parameter-shift differentiation
# mode, then there must be no observable following it.
continue
B = BB.heisenberg_tr(dev_wires) @ B
B_inv = B_inv @ BB.heisenberg_tr(dev_wires, inverse=True)
Z = B @ Z @ B_inv # conjugation
new_measurements = list(tape.measurements)
constants = []
# transform the descendant observables into their derivatives using Z
transformed_obs_idx = []
for mp in observable_descendents:
obs = mp if mp.obs is None else mp.obs
# get the index of the descendent observable
# pylint:disable=undefined-loop-variable
for obs_idx, tape_obs in enumerate(tape.observables):
if obs is tape_obs:
break
transformed_obs_idx.append(obs_idx)
transformed_obs = _transform_observable(obs, Z, dev_wires)
A = transformed_obs.parameters[0]
constant = None
# Check if the transformed observable corresponds to a constant term.
if len(A.nonzero()[0]) == 1:
if A.ndim == 2 and A[0, 0] != 0:
constant = A[0, 0]
elif A.ndim == 1 and A[0] != 0:
constant = A[0]
constants.append(constant)
new_measurements[obs_idx] = qml.expval(op=_transform_observable(obs, Z, dev_wires))
g_tape = qml.tape.QuantumScript(
tape.operations,
new_measurements,
shots=tape.shots,
trainable_params=tape.trainable_params,
)
if not any(i is None for i in constants):
# Check if *all* transformed observables corresponds to a constant
# term. If this is the case for all transformed observables on the
# tape, then <psi|A|psi> = A<psi|psi> = A, and we can avoid the
# device execution.
shapes[-1] = 0
obs_indices.append(transformed_obs_idx)
gradient_values.append(constants)
continue
gradient_tapes.append(g_tape)
obs_indices.append(transformed_obs_idx)
gradient_values.append(None)
def processing_fn(results):
grads = []
start = 0
if not results:
results = [np.squeeze(np.zeros([tape.output_dim]))]
interface = qml.math.get_interface(results[0])
iterator = enumerate(zip(shapes, gradient_values, obs_indices))
for i, (shape, grad_value, obs_ind) in iterator:
if shape == 0:
# parameter has zero gradient
isscalar = qml.math.ndim(results[0]) == 0
g = qml.math.zeros_like(qml.math.atleast_1d(results[0]), like=interface)
if grad_value:
g = qml.math.scatter_element_add(g, obs_ind, grad_value, like=interface)
grads.append(g[0] if isscalar else g)
continue
obs_result = results[start : start + shape]
start = start + shape
# compute the linear combination of results and coefficients
isscalar = qml.math.ndim(obs_result[0]) == 0
obs_result = qml.math.stack(qml.math.atleast_1d(obs_result[0]))
g = qml.math.zeros_like(obs_result, like=interface)
if qml.math.get_interface(g) not in ("tensorflow", "autograd"):
obs_ind = (obs_ind,)
g = qml.math.scatter_element_add(g, obs_ind, obs_result[obs_ind], like=interface)
grads.append(g[0] if isscalar else g)
# The following is for backwards compatibility; currently,
# the device stacks multiple measurement arrays, even if not the same
# size, resulting in a ragged array.
# In the future, we might want to change this so that only tuples
# of arrays are returned.
for i, g in enumerate(grads):
g = qml.math.convert_like(g, results[0])
if hasattr(g, "dtype") and g.dtype is np.dtype("object"):
grads[i] = qml.math.hstack(g)
return qml.math.T(qml.math.stack(grads))
return gradient_tapes, processing_fn
def _expand_transform_param_shift_cv(
tape: QuantumScript,
dev,
argnum=None,
shifts=None,
gradient_recipes=None,
fallback_fn=finite_diff,
f0=None,
force_order2=False,
) -> tuple[QuantumScriptBatch, PostprocessingFn]:
"""Expand function to be applied before parameter shift CV."""
expanded_tape = expand_invalid_trainable(tape)
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]
return [expanded_tape], null_postprocessing
[docs]@partial(
transform,
expand_transform=_expand_transform_param_shift_cv,
classical_cotransform=_contract_qjac_with_cjac,
final_transform=True,
)
def param_shift_cv(
tape: QuantumScript,
dev,
argnum=None,
shifts=None,
gradient_recipes=None,
fallback_fn=finite_diff,
f0=None,
force_order2=False,
) -> tuple[QuantumScriptBatch, PostprocessingFn]:
r"""Transform a continuous-variable QNode to compute the parameter-shift gradient of all gate
parameters with respect to its inputs.
Args:
tape (QNode or QuantumTape): quantum circuit to differentiate
dev (pennylane.Device): device the parameter-shift method is to be computed on
argnum (int or list[int] or None): Trainable parameter indices to differentiate
with respect to. If not provided, the derivative with respect to all
trainable indices are returned.
shifts (list[tuple[int or float]]): List containing tuples of shift values.
If provided, one tuple of shifts should be given per trainable parameter
and the tuple should match the number of frequencies for that parameter.
If unspecified, equidistant shifts are assumed.
gradient_recipes (tuple(list[list[float]] or None)): List of gradient recipes
for the parameter-shift method. One gradient recipe must be provided
per trainable parameter.
This is a tuple with one nested list per parameter. For
parameter :math:`\phi_k`, the nested list contains elements of the form
:math:`[c_i, a_i, s_i]` where :math:`i` is the index of the
term, resulting in a gradient recipe of
.. math:: \frac{\partial}{\partial\phi_k}f = \sum_{i} c_i f(a_i \phi_k + s_i).
If ``None``, the default gradient recipe containing the two terms
:math:`[c_0, a_0, s_0]=[1/2, 1, \pi/2]` and :math:`[c_1, a_1,
s_1]=[-1/2, 1, -\pi/2]` is assumed for every parameter.
fallback_fn (None or Callable): a fallback gradient function to use for
any parameters that do not support the parameter-shift rule.
f0 (tensor_like[float] or None): Output of the evaluated input tape. If provided,
and the gradient recipe contains an unshifted term, this value is used,
saving a quantum evaluation.
force_order2 (bool): if True, use the order-2 method even if not necessary
Returns:
qnode (QNode) or tuple[List[QuantumTape], function]:
The transformed circuit as described in :func:`qml.transform <pennylane.transform>`. Executing this circuit
will provide the Jacobian in the form of a tensor, a tuple, or a nested tuple depending upon the nesting
structure of measurements in the original circuit.
This transform supports analytic gradients of Gaussian CV operations using
the parameter-shift rule. This gradient method returns *exact* gradients,
and can be computed directly on quantum hardware.
Analytic gradients of photonic circuits that satisfy
the following constraints with regards to measurements are supported:
* Expectation values are restricted to observables that are first- and
second-order in :math:`\hat{x}` and :math:`\hat{p}` only.
This includes :class:`~.X`, :class:`~.P`, :class:`~.QuadOperator`,
:class:`~.PolyXP`, and :class:`~.NumberOperator`.
For second-order observables, the device **must support** :class:`~.PolyXP`.
* Variances are restricted to observables that are first-order
in :math:`\hat{x}` and :math:`\hat{p}` only. This includes :class:`~.X`, :class:`~.P`,
:class:`~.QuadOperator`, and *some* parameter values of :class:`~.PolyXP`.
The device **must support** :class:`~.PolyXP`.
.. warning::
Fock state probabilities (tapes that return :func:`~pennylane.probs` or
expectation values of :class:`~.FockStateProjector`) are not supported.
In addition, the operations must fulfill the following requirements:
* Only Gaussian operations are differentiable.
* Non-differentiable Fock states and Fock operations may *precede* all differentiable Gaussian,
operations. For example, the following is permissible:
.. code-block:: python
@qml.qnode(dev)
def circuit(weights):
# Non-differentiable Fock operations
qml.FockState(np.array(2, requires_grad=False), wires=0)
qml.Kerr(np.array(0.654, requires_grad=False), wires=1)
# differentiable Gaussian operations
qml.Displacement(weights[0], weights[1], wires=0)
qml.Beamsplitter(weights[2], weights[3], wires=[0, 1])
return qml.expval(qml.NumberOperator(0))
* If a Fock operation succeeds a Gaussian operation, the Fock operation must
not contribute to any measurements. For example, the following is allowed:
.. code-block:: python
@qml.qnode(dev)
def circuit(weights):
qml.Displacement(weights[0], weights[1], wires=0)
qml.Beamsplitter(weights[2], weights[3], wires=[0, 1])
qml.Kerr(np.array(0.654, requires_grad=False), wires=1) # there is no measurement on wire 1
return qml.expval(qml.NumberOperator(0))
If any of the above constraints are not followed, the tape cannot be differentiated
via the CV parameter-shift rule. Please use numerical differentiation instead.
**Example**
This transform can be registered directly as the quantum gradient transform
to use during autodifferentiation:
>>> dev = qml.device("default.gaussian", wires=2)
>>> @qml.qnode(dev, diff_method="parameter-shift")
... def circuit(params):
... qml.Squeezing(params[0], params[1], wires=[0])
... qml.Squeezing(params[2], params[3], wires=[0])
... return qml.expval(qml.NumberOperator(0))
>>> params = np.array([0.1, 0.2, 0.3, 0.4], requires_grad=True)
>>> qml.jacobian(circuit)(params)
array([ 0.87516064, 0.01273285, 0.88334834, -0.01273285])
.. details::
:title: Usage Details
This gradient transform can be applied directly to :class:`QNode <pennylane.QNode>` objects.
However, for performance reasons, we recommend providing the gradient transform as the ``diff_method`` argument
of the QNode decorator, and differentiating with your preferred machine learning framework.
>>> @qml.qnode(dev)
... def circuit(params):
... qml.Squeezing(params[0], params[1], wires=[0])
... qml.Squeezing(params[2], params[3], wires=[0])
... return qml.expval(qml.NumberOperator(0))
>>> params = np.array([0.1, 0.2, 0.3, 0.4], requires_grad=True)
>>> qml.gradients.param_shift_cv(circuit, dev)(params)
tensor([[ 0.87516064, 0.01273285, 0.88334834, -0.01273285]], requires_grad=True)
This quantum gradient transform can also be applied to low-level
:class:`~.QuantumTape` objects. This will result in no implicit quantum
device evaluation. Instead, the processed tapes, and post-processing
function, which together define the gradient are directly returned:
>>> r0, phi0, r1, phi1 = [0.4, -0.3, -0.7, 0.2]
>>> ops = [qml.Squeezing(r0, phi0, wires=0), qml.Squeezing(r1, phi1, wires=0)]
>>> tape = qml.tape.QuantumTape(ops, [qml.expval(qml.NumberOperator(0))])
>>> gradient_tapes, fn = qml.gradients.param_shift_cv(tape, dev)
>>> gradient_tapes
[<QuantumTape: wires=[0], params=4>,
<QuantumTape: wires=[0], params=4>,
<QuantumTape: wires=[0], params=4>,
<QuantumTape: wires=[0], params=4>]
This can be useful if the underlying circuits representing the gradient
computation need to be analyzed.
The output tapes can then be evaluated and post-processed to retrieve
the gradient:
>>> dev = qml.device("default.gaussian", wires=2)
>>> fn(qml.execute(gradient_tapes, dev, None))
(-0.32487113372219933,
-0.4054074025310772,
-0.8704985300843778,
0.4054074025310775)
"""
if len(tape.measurements) > 1:
raise ValueError(
"Computing the gradient of CV circuits that return more than one measurement is not possible."
)
# perform gradient method validation
if any(isinstance(m, StateMP) for m in tape.measurements):
raise ValueError(
"Computing the gradient of circuits that return the state is not supported."
)
method = "analytic" if fallback_fn is None else "best"
trainable_params = choose_trainable_params(tape, argnum)
method_map = _gradient_analysis_and_validation_cv(tape, method, trainable_params)
if argnum is None and not tape.trainable_params:
return _no_trainable_grad(tape)
gradient_tapes = []
shapes = []
fns = []
def _update(data):
"""Utility function to update the list of gradient tapes,
the corresponding number of gradient tapes, and the processing functions"""
gradient_tapes.extend(data[0])
shapes.append(len(data[0]))
fns.append(data[1])
if all(g == "0" for g in method_map.values()):
return [], lambda _: np.zeros([tape.output_dim, len(tape.trainable_params)])
var_present = any(isinstance(m, VarianceMP) for m in tape.measurements)
unsupported_params = []
first_order_params = []
second_order_params = []
for idx, g in method_map.items():
if g == "F":
unsupported_params.append(idx)
elif g == "A":
first_order_params.append(idx)
elif g == "A2":
second_order_params.append(idx)
if force_order2:
# all analytic parameters should be computed using the second-order method
second_order_params += first_order_params
first_order_params = []
if "PolyXP" not in dev.observables and (second_order_params or var_present):
warnings.warn(
f"The device {dev.short_name} does not support "
"the PolyXP observable. The analytic parameter-shift cannot be used for "
"second-order observables; falling back to finite-differences.",
UserWarning,
)
if var_present:
unsupported_params += first_order_params
first_order_params = []
unsupported_params += second_order_params
second_order_params = []
# If there are unsupported operations, call the fallback gradient function
if unsupported_params:
_update(fallback_fn(tape, argnum=unsupported_params))
# collect all the analytic parameters
argnum = first_order_params + second_order_params
if not argnum:
# No analytic parameters. Return the existing fallback tapes/fn
return gradient_tapes, fns[-1]
gradient_recipes = gradient_recipes or [None] * len(argnum)
if var_present:
_update(var_param_shift(tape, dev.wires, argnum, shifts, gradient_recipes, f0))
else:
# Only expectation values were specified
if first_order_params:
_update(expval_param_shift(tape, first_order_params, shifts, gradient_recipes, f0))
if second_order_params:
_update(
second_order_param_shift(
tape, dev.wires, second_order_params, shifts, gradient_recipes
)
)
def processing_fn(results):
start = 0
grads = []
for s, f in zip(shapes, fns):
grads.append(f(results[start : start + s]))
start += s
# For expval param shift with multiple params
if isinstance(grads[0], tuple):
grads = [qml.math.stack(g) for g in grads]
jacobian = sum(grads)
if jacobian.shape != ():
if jacobian.shape[0] == 1:
jacobian = jacobian[0]
if len(argnum) > 1:
jacobian = tuple(j for j in jacobian)
return jacobian
return gradient_tapes, processing_fn
_modules/pennylane/gradients/parameter_shift_cv
Download Python script
Download Notebook
View on GitHub