# Copyright 2018-2021 Xanadu Quantum Technologies Inc.

# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

# Unless required by applicable law or agreed to in writing, software
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
"""
This module contains functions for computing the Jacobian vector product
of tapes.
"""
from collections.abc import Sequence

import numpy as np

import pennylane as qml
from pennylane._device import _get_num_copies
from pennylane.measurements import ProbabilityMP

def _convert(jac, tangent):
"""Utility to convert and cast the jacobian as tangent."""
if isinstance(jac, tuple):
jac_new = []
for j in jac:
j_ = qml.math.convert_like(j, tangent)
j_ = qml.math.cast(j_, tangent.dtype)
jac_new.append(j_)
jac = tuple(jac_new)
else:
jac = qml.math.convert_like(jac, tangent)
jac = qml.math.cast(jac, tangent.dtype)
return jac

[docs]def compute_jvp_single(tangent, jac):
"""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**

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)
np.array(0.2)

2. For a single parameter and a single measurment with shape (e.g. probs):

.. code-block:: pycon

>>> tangent = np.array([2.0])
>>> jac = np.array([0.3, 0.4])
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)])
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])])
np.array([0.2, 0.5])

"""
if jac is None:
return None

tangent = qml.math.stack(tangent)
jac = _convert(jac, tangent)

# Single param
if not isinstance(jac, tuple):
# No trainable parameters
if jac.shape == (0,):
res = qml.math.zeros((1, 0))
return res

tangent = qml.math.reshape(tangent, (1,))

# No dimension e.g. expval
if jac.shape == ():
jac = qml.math.reshape(jac, (1,))

# With dimension e.g. probs
else:
jac = qml.math.reshape(jac, (1, -1))
res = qml.math.tensordot(jac, tangent, [, ])
# Multiple params
else:
if len(jac) == 0:
res = qml.math.zeros((1, 0))
return res

jac = qml.math.stack(jac)

# No dimension e.g. expval
if jac.shape == ():
res = qml.math.tensordot(jac, tangent, 1)

# With dimension e.g. probs
else:
res = qml.math.tensordot(jac, tangent, [, ])
return res

[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])])
(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])]),])
(np.array([1.1]), np.array([0.8, 2.1]))
"""
if jac is None:
return None
res = tuple(compute_jvp_single(tangent, j) for j in jac)
return res

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.
the tape
shots (None, int, list[int]): The device shots that will be used to
execute the tapes outputted by this
gradient_kwargs (dict): dictionary of keyword arguments to pass when

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

qml.enable_return()

x = jax.numpy.array([[0.1, 0.2, 0.3],
[0.4, 0.5, 0.6]])

with qml.tape.QuantumTape() as tape:
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)
qml.expval(qml.PauliZ(0))
qml.probs(wires=1)

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)]

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.batch_execute(jvp_tapes))
>>> jvp
(Array(-0.62073976, dtype=float32), Array([-0.3259707 ,  0.32597077], dtype=float32))
"""
num_params = len(tape.trainable_params)
num_measurements = len(tape.measurements)

if num_params == 0:
# The tape has no trainable parameters; the JVP
# is simply none.
return [], lambda _, num=None: None

multi_m = num_measurements > 1

try:
if qml.math.allclose(qml.math.stack(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
if not multi_m:
# TODO: Update shape for CV variables and for qutrit simulations
res = _single_measurement_zero(tape.measurements, tangent)
else:
# TODO: Update shape for CV variables and for qutrit simulations
res = [_single_measurement_zero(m, tangent) for m in tape.measurements]
res = tuple(res)
return res

return [], func
except (AttributeError, TypeError):
pass

def processing_fn(results):
# postprocess results to compute the Jacobian
jac = fn(results)
shot_vector = isinstance(shots, Sequence)

# Jacobian without shot vectors
if not shot_vector:
if multi_m:
return compute_jvp_multi(tangent, jac)
return compute_jvp_single(tangent, jac)

# The jacobian is calculated for shot vectors
len_shot_vec = _get_num_copies(shots)
jvps = []
if multi_m:
for i in range(len_shot_vec):
jvps.append(compute_jvp_multi(tangent, jac[i]))
else:
for i in range(len_shot_vec):
jvps.append(compute_jvp_single(tangent, jac[i]))

return tuple(jvps)

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.
the tapes
shots (None, int, list[int]): The device shots that will be used to
execute the tapes outputted by this
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

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
qml.enable_return()
x = jax.numpy.array([[0.1, 0.2, 0.3],
[0.4, 0.5, 0.6]])

def ansatz(x):
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)

with qml.tape.QuantumTape() as tape1:
ansatz(x)
qml.expval(qml.PauliZ(0))
qml.probs(wires=1)

with qml.tape.QuantumTape() as tape2:
ansatz(x)
qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))

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:

>>> dev = qml.device("default.qubit", wires=2)
>>> jvps = fn(dev.batch_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
reshape_info = []
processing_fns = []

# Loop through the tapes and dys vector
for tape, tangent in zip(tapes, tangents):

reshape_info.append(len(g_tapes))
processing_fns.append(fn)

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 jvps

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


Using PennyLane

Development

API

Internals