Source code for pennylane.gradients.jvp

# 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 Jacobian vector product
of tapes.
"""
import numpy as np

import pennylane as qml
from pennylane.measurements import ProbabilityMP


[docs]def compute_jvp_single(tangent, jac): r"""Convenience function to compute the Jacobian vector product for a given tangent vector and a Jacobian for a single measurement tape. Args: tangent (list, tensor_like): tangent vector jac (tensor_like, tuple): Jacobian matrix Returns: tensor_like: the Jacobian vector product **Examples** We start with a number of examples. A more complete, technical description is given further below. 1. For a single parameter and a single measurement without shape (e.g. ``expval``, ``var``): .. code-block:: pycon >>> tangent = np.array([1.0]) >>> jac = np.array(0.2) >>> qml.gradients.compute_jvp_single(tangent, jac) np.array(0.2) 2. For a single parameter and a single measurement with shape (e.g. ``probs``): .. code-block:: pycon >>> tangent = np.array([2.0]) >>> jac = np.array([0.3, 0.4]) >>> qml.gradients.compute_jvp_single(tangent, jac) np.array([0.6, 0.8]) 3. For multiple parameters (in this case 2 parameters) and a single measurement without shape (e.g. ``expval``, ``var``): .. code-block:: pycon >>> tangent = np.array([1.0, 2.0]) >>> jac = tuple([np.array(0.1), np.array(0.2)]) >>> qml.gradients.compute_jvp_single(tangent, jac) np.array(0.5) 4. For multiple parameters (in this case 2 parameters) and a single measurement with shape (e.g. ``probs``): .. code-block:: pycon >>> tangent = np.array([1.0, 0.5]) >>> jac = tuple([np.array([0.1, 0.3]), np.array([0.2, 0.4])]) >>> qml.gradients.compute_jvp_single(tangent, jac) np.array([0.2, 0.5]) .. details:: :title: Technical description :href: technical-description There are multiple case distinctions in this function, for particular examples see above. - The JVP may be for one **(A)** or multiple **(B)** parameters. We call the number of parameters ``k`` - The number ``R`` of tape return type dimensions may be between 0 and 3. We call the return type dimensions ``r_j`` - Each parameter may have an arbitrary number ``L_i>=0`` of dimensions In the following, ``(a, b)`` denotes a tensor_like of shape ``(a, b)`` and ``[(a,), (b,)]`` / ``((a,), (b,))`` denotes a ``list`` / ``tuple`` of tensors with the indicated shapes, respectively. Ignore the case of no trainable parameters, as it is filtered out in advance. For scenario **(A)**, the input shapes can be in .. list-table:: :widths: 30 40 30 :header-rows: 1 * - ``tangent`` shape - ``jac`` shape - Comment * - ``(1,)`` or ``[()]`` or ``(())`` - ``()`` - scalar return, scalar parameter * - ``(1,)`` or ``[()]`` or ``(())`` - ``(r_1,..,r_R)`` - tensor return, scalar parameter * - ``[(l_1,..,l_{L_1})]`` [1] - ``(l_1,..,l_{L_1})`` - scalar return, tensor parameter * - ``[(l_1,..,l_{L_1})]`` [1] - ``(r_1,..,r_R, l_1,..,l_{L_1})`` - tensor return, tensor parameter [1] Note that intuitively, ``tangent`` could be allowed to be a tensor of shape ``(l_1,..,l_{L_1})`` without an outer list. However, this is excluded in order to allow for the distinction from scenario **(B)**. Internally, this input shape for ``tangent`` never occurs for scenario **(A)**. In this scenario, the tangent is reshaped into a one-dimensional tensor with shape ``(tangent_size,)`` and the Jacobian is reshaped to have the dimensions ``(r_1, ... r_R, tangent_size)``. This is followed by a ``tensordot`` contraction over the ``tangent_size`` axis of both tensors. For scenario **(B)**, the input shapes can be in .. list-table:: :widths: 30 40 30 :header-rows: 1 * - ``tangent`` shape - ``jac`` shape - Comment * - ``(k,)`` or ``[(),..,()]`` or ``((),..,())`` - ``((),..,())`` (length ``k``) - scalar return, ``k`` scalar parameters * - ``(k,)`` or ``[(),..,()]`` or ``((),..,())`` - ``((r_1,..,r_R),..,(r_1,..,r_R))`` [1] - tensor return, ``k`` scalar parameters * - ``[(l_1,..,l_{L_1}),..,(l_1,..,l_{L_k})]`` - ``((l_1,..,l_{L_1}),..,(l_1,..,l_{L_k}))`` - scalar return, ``k`` tensor parameters * - ``[(l_1,..,l_{L_1}),..,(l_1,..,l_{L_k})]`` - ``((r_1,..,r_R, l_1,..,l_{L_1}),..,(r_1,..,r_R, l_1,..,l_{L_k}))`` [1] - tensor return, ``k`` tensor parameters [1] Note that the return type dimensions ``(r_1,..,r_R)`` are the same for all entries of ``jac``, whereas the dimensions of the entries in ``tanget``, and the according dimensions ``(l_1,..,l_{L_k})`` of the ``jac`` entries may differ. In this scenario, another case separation is used: If any of the parameters is a tensor (i.e. not a scalar), all tangent entries are reshaped into one-dimensional tensors with shapes ``(tangent_size_i,)`` and then stacked into one one-dimensional tensor. If there are no tensor parameters, the tangent is just stacked and reshaped. The Jacobians are reshaped to have the dimensions ``(r_1, ... r_R, tangent_size_i)`` and then are concatenated along their last (potentially mismatching) axis. This is followed by a tensordot contraction over the axes of size :math:`\sum_i` ``tangent_size_i``. """ if jac is None: return None single_param = not isinstance(jac, tuple) if (single_param and jac.shape == (0,)) or (not single_param and len(jac) == 0): # No trainable parameters return qml.math.zeros((1, 0)) if single_param: tangent = qml.math.stack(tangent) first_tangent_ndim = len(tangent.shape[1:]) tangent = qml.math.flatten(tangent) tangent_size = tangent.shape[0] shape = jac.shape new_shape = shape[: len(shape) - first_tangent_ndim] + (tangent_size,) jac = qml.math.cast(qml.math.convert_like(jac, tangent), tangent.dtype) jac = qml.math.reshape(jac, new_shape) return qml.math.tensordot(jac, tangent, [[-1], [0]]) tangent_ndims = [getattr(t, "ndim", 0) for t in tangent] if isinstance(tangent, (tuple, list)) and any(ndim > 0 for ndim in tangent_ndims): # At least one tangent entry is not a scalar, requiring us to flatten them and hstack tangent = [qml.math.flatten(t) for t in tangent] tangent_sizes = [t.shape[0] for t in tangent] tangent = qml.math.hstack(tangent) else: # Only scalar tangent entries, no flattening required and we may use stack tangent_sizes = [1] * len(tangent) tangent = qml.math.stack(tangent) jac_shapes = [j.shape for j in jac] new_shapes = [ shape[: len(shape) - t_ndim] + (tsize,) for shape, t_ndim, tsize in zip(jac_shapes, tangent_ndims, tangent_sizes) ] jac = qml.math.concatenate([qml.math.reshape(j, s) for j, s in zip(jac, new_shapes)], axis=-1) jac = qml.math.cast(qml.math.convert_like(jac, tangent), tangent.dtype) return qml.math.tensordot(jac, tangent, [[-1], [0]])
[docs]def compute_jvp_multi(tangent, jac): """Convenience function to compute the Jacobian-vector product for a given vector of gradient outputs and a Jacobian for a tape with multiple measurements. Args: tangent (tensor_like, list): tangent vector jac (tensor_like, tuple): Jacobian matrix Returns: tensor_like: the Jacobian-vector product **Examples** 1. For a single parameter and multiple measurements (one without shape and one with shape, e.g. expval and probs): .. code-block:: pycon >>> tangent = np.array([2.0]) >>> jac = tuple([np.array([0.3]), np.array([0.2, 0.5])]) >>> qml.gradients.compute_jvp_multi(tangent, jac) (np.array([0.6]), np.array([0.4, 1. ])) 2. For multiple parameters (in this case 2 parameters) and multiple measurements (one without shape and one with shape, e.g. expval and probs): .. code-block:: pycon >>> tangent = np.array([1.0, 2.0]) >>> jac = tuple([tuple([np.array([0.3]), np.array([0.4])]), tuple([np.array([0.2, 0.5]), np.array([0.3, 0.8])]),]) >>> qml.gradients.compute_jvp_multi(tangent, jac) (np.array([1.1]), np.array([0.8, 2.1])) """ if jac is None: return None return tuple(compute_jvp_single(tangent, j) for j in jac)
[docs]def jvp(tape, tangent, gradient_fn, gradient_kwargs=None): r"""Generate the gradient tapes and processing function required to compute the Jacobian vector product of a tape. This function only works with the new return type system on. Args: tape (.QuantumTape): quantum tape to differentiate tangent (tensor_like, list): Gradient-output vector. Must have shape matching the number of trainable parameters. gradient_fn (callable): the gradient transform to use to differentiate the tape gradient_kwargs (dict): dictionary of keyword arguments to pass when determining the gradients of tapes Returns: tensor_like or tuple or None: Jacobian vector product. Returns None if the tape has no trainable parameters. **Example** Consider the following quantum tape with Jax parameters: .. code-block:: python import jax x = jax.numpy.array([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]) ops = [ qml.RX(x[0, 0], wires=0), qml.RY(x[0, 1], wires=1), qml.RZ(x[0, 2], wires=0), qml.CNOT(wires=[0, 1]), qml.RX(x[1, 0], wires=1), qml.RY(x[1, 1], wires=0), qml.RZ(x[1, 2], wires=1) ] measurements = [qml.expval(qml.PauliZ(0)), qml.probs(wires=1)] tape = qml.tape.QuantumTape(ops, measurements) We can use the ``jvp`` function to compute the Jacobian vector product, given a tangent vector ``tangent``: >>> tangent = [jax.numpy.array(1.0), jax.numpy.array(1.0), jax.numpy.array(1.0), jax.numpy.array(1.0), jax.numpy.array(1.0), jax.numpy.array(1.0)] >>> jvp_tapes, fn = qml.gradients.jvp(tape, tangent, qml.gradients.param_shift) Note that ``tangent`` has six elements, matching the parameter dimension of the tape. Executing the JVP tapes, and applying the processing function: >>> dev = qml.device("default.qubit", wires=2) >>> jvp = fn(dev.execute(jvp_tapes)) >>> jvp (Array(-0.62073976, dtype=float32), Array([-0.3259707 , 0.32597077], dtype=float32)) """ if len(tape.trainable_params) == 0: # The tape has no trainable parameters; the JVP # is simply none. def zero_vjp(_): res = tuple(np.zeros(mp.shape(None, tape.shots)) for mp in tape.measurements) return res[0] if len(tape.measurements) == 1 else res return tuple(), zero_vjp multi_m = len(tape.measurements) > 1 try: # if qml.math.allclose(qml.math.stack(tangent), 0): if qml.math.allclose(tangent, 0): # If the tangent vector is zero, then the # corresponding element of the JVP will be zero, # and we can avoid a quantum computation. def func(_): # pylint: disable=unused-argument # TODO: Update shape for CV variables and for qutrit simulations res = tuple(_single_measurement_zero(m, tangent) for m in tape.measurements) if not multi_m: res = res[0] return res return [], func except (AttributeError, TypeError): pass gradient_kwargs = gradient_kwargs or {} gradient_tapes, fn = gradient_fn(tape, **gradient_kwargs) def processing_fn(results): # postprocess results to compute the Jacobian jac = fn(results) _jvp_fn = compute_jvp_multi if multi_m else compute_jvp_single # Jacobian without shot vectors if not tape.shots.has_partitioned_shots: return _jvp_fn(tangent, jac) # The jacobian is calculated for shot vectors return tuple(_jvp_fn(tangent, jac[i]) for i in range(tape.shots.num_copies)) return gradient_tapes, processing_fn
[docs]def batch_jvp(tapes, tangents, gradient_fn, reduction="append", gradient_kwargs=None): r"""Generate the gradient tapes and processing function required to compute the Jacobian vector products of a batch of tapes. Args: tapes (Sequence[.QuantumTape]): sequence of quantum tapes to differentiate tangents (Sequence[tensor_like]): Sequence of gradient-output vectors ``dy``. Must be the same length as ``tapes``. Each ``dy`` tensor should have shape matching the output shape of the corresponding tape. gradient_fn (callable): the gradient transform to use to differentiate the tapes reduction (str): Determines how the Jacobian-vector products are returned. If ``append``, then the output of the function will be of the form ``List[tensor_like]``, with each element corresponding to the JVP of each input tape. If ``extend``, then the output JVPs will be concatenated. gradient_kwargs (dict): dictionary of keyword arguments to pass when determining the gradients of tapes Returns: List[tensor_like or None]: list of Jacobian vector products. ``None`` elements corresponds to tapes with no trainable parameters. **Example** .. code-block:: python import jax x = jax.numpy.array([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]) ops = [ qml.RX(x[0, 0], wires=0), qml.RY(x[0, 1], wires=1), qml.RZ(x[0, 2], wires=0), qml.CNOT(wires=[0, 1]), qml.RX(x[1, 0], wires=1), qml.RY(x[1, 1], wires=0), qml.RZ(x[1, 2], wires=1) ] measurements1 = [qml.expval(qml.PauliZ(0)), qml.probs(wires=1)] tape1 = qml.tape.QuantumTape(ops, measurements1) measurements2 = [qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))] tape2 = qml.tape.QuantumTape(ops, measurements2) tapes = [tape1, tape2] Both tapes share the same circuit ansatz, but have different measurement outputs. We can use the ``batch_jvp`` function to compute the Jacobian vector product, given a list of tangents ``tangent``: >>> tangent_0 = [jax.numpy.array(1.0), jax.numpy.array(1.0), jax.numpy.array(1.0), jax.numpy.array(1.0), jax.numpy.array(1.0), jax.numpy.array(1.0)] >>> tangent_1 = [jax.numpy.array(1.0), jax.numpy.array(1.0), jax.numpy.array(1.0), jax.numpy.array(1.0), jax.numpy.array(1.0), jax.numpy.array(1.0)] >>> tangents = [tangent_0, tangent_1] Note that each ``tangents`` has shape matching the parameter dimension of the tape. Executing the JVP tapes, and applying the processing function: >>> jvp_tapes, fn = qml.gradients.batch_jvp(tapes, tangents, qml.gradients.param_shift) >>> dev = qml.device("default.qubit", wires=2) >>> jvps = fn(dev.execute(jvp_tapes)) >>> jvps ((Array(-0.62073976, dtype=float32), Array([-0.3259707 , 0.32597077], dtype=float32)), Array(-0.6900841, dtype=float32)) We have two JVPs; one per tape. Each one corresponds to the shape of the output of their respective tape. """ # pylint: disable=too-many-arguments gradient_kwargs = gradient_kwargs or {} reshape_info = [] gradient_tapes = [] processing_fns = [] # Loop through the tapes and dys vector for tape, tangent in zip(tapes, tangents): g_tapes, fn = jvp(tape, tangent, gradient_fn, gradient_kwargs) reshape_info.append(len(g_tapes)) processing_fns.append(fn) gradient_tapes.extend(g_tapes) def processing_fn(results): jvps = [] start = 0 for t_idx in range(len(tapes)): # extract the correct results from the flat list res_len = reshape_info[t_idx] res_t = results[start : start + res_len] start += res_len # postprocess results to compute the JVP jvp_ = processing_fns[t_idx](res_t) if jvp_ is None: if reduction == "append": jvps.append(None) continue if isinstance(reduction, str): getattr(jvps, reduction)(jvp_) elif callable(reduction): reduction(jvps, jvp_) return tuple(jvps) return gradient_tapes, processing_fn
def _single_measurement_zero(m, tangent): """Aux function to create a zero tensor from a measurement.""" dim = 2 ** len(m.wires) if isinstance(m, ProbabilityMP) else () res = qml.math.convert_like(np.zeros(dim), tangent) res = qml.math.cast_like(res, tangent) return res