# Source code for pennylane.transforms.classical_jacobian

# 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
"""
Contains the classical Jacobian transform.
"""
# pylint: disable=import-outside-toplevel
import pennylane as qml
from pennylane import numpy as np

[docs]def classical_jacobian(qnode, argnum=None, expand_fn=None, trainable_only=True):
r"""Returns a function to extract the Jacobian
matrix of the classical part of a QNode.

This transform allows the classical dependence between the QNode
arguments and the quantum gate arguments to be extracted.

Args:
qnode (pennylane.QNode): QNode to compute the (classical) Jacobian of
argnum (int or Sequence[int]): indices of QNode arguments with respect to which
the (classical) Jacobian is computed
expand_fn (None or function): an expansion function (if required) to be applied to the
QNode quantum tape before the classical Jacobian is computed

Returns:
function: Function which accepts the same arguments as the QNode.
When called, this function will return the Jacobian of the QNode
gate arguments with respect to the QNode arguments indexed by argnum.

**Example**

Consider the following QNode:

>>> @qml.qnode(dev)
... def circuit(weights):
...     qml.RX(weights, wires=0)
...     qml.RY(0.2 * weights, wires=1)
...     qml.RY(2.5, wires=0)
...     qml.RZ(weights ** 2, wires=1)
...     qml.RX(weights, wires=1)
...     return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))

We can use this transform to extract the relationship :math:f: \mathbb{R}^n \rightarrow
\mathbb{R}^m between the input QNode arguments :math:w and the gate arguments :math:g, for
a given value of the QNode arguments:

>>> cjac_fn = qml.transforms.classical_jacobian(circuit)
>>> weights = np.array([1., 1., 0.6], requires_grad=True)
>>> cjac = cjac_fn(weights)
>>> print(cjac)
[[1.  0.  0. ]
[0.2 0.  0. ]
[0.  0.  0. ]
[0.  1.2 0. ]
[0.  0.  1. ]]

The returned Jacobian has rows corresponding to gate arguments, and columns
corresponding to QNode arguments; that is,

.. math:: J_{ij} = \frac{\partial}{\partial g_i} f(w_j).

We can see that:

- The zeroth element of weights is repeated on the first two gates generated by the QNode.

- The third row consisting of all zeros indicates that the third gate RY(2.5) does not
depend on the weights.

- The quadratic dependence of the fourth gate argument yields :math:2\cdot 0.6=1.2.

.. note::

The QNode is constructed during this operation.

For a QNode with multiple QNode arguments, the arguments with respect to which the
Jacobian is computed can be controlled with the argnum keyword argument.
The output and its format depend on the backend:

.. list-table:: Output format of classical_jacobian
:widths: 15 25 25 35

* - Interface
- argnum=None
- type(argnum)=int
- type(argnum) = Sequence[int]
* - 'autograd'
- tuple(array) 
- array
- tuple(array)
* - 'jax'
- array 
- array
- tuple(array)
* - 'tf'
- tuple(array)
- array
- tuple(array)
* - 'torch'
- tuple(array)
- array
- tuple(array)

 If there only is one trainable QNode argument, the tuple is unpacked to a
single array, as is the case for :func:.jacobian.

 For JAX, argnum=None defaults to argnum=0 in contrast to all other
interfaces. This means that only the classical Jacobian with respect to the first
QNode argument is computed if no argnum is provided.

**Example with argnum**

>>> @qml.qnode(dev)
... def circuit(x, y, z):
...     qml.RX(qml.math.sin(x), wires=0)
...     qml.CNOT(wires=[0, 1])
...     qml.RY(y ** 2, wires=1)
...     qml.RZ(1 / z, wires=1)
...     return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))
>>> jac_fn = qml.transforms.classical_jacobian(circuit, argnum=[1, 2])
>>> x, y, z = np.array([0.1, -2.5, 0.71])
>>> jac_fn(x, y, z)
(array([-0., -5., -0.]), array([-0.        , -0.        , -1.98373339]))

Only the Jacobians with respect to the arguments x and y were computed, and
returned as a tuple of arrays.
"""

def classical_preprocessing(*args, **kwargs):
"""Returns the trainable gate parameters for a given QNode input."""
kwargs.pop("shots", None)
kwargs.pop("argnums", None)
qnode.construct(args, kwargs)
tape = qnode.qtape

if expand_fn is not None:
tape = expand_fn(tape)
return qml.math.stack(tape.get_parameters(trainable_only=trainable_only))

wrapper_argnum = argnum if argnum is not None else None

def qnode_wrapper(*args, **kwargs):  # pylint: disable=inconsistent-return-statements
old_interface = qnode.interface

if old_interface == "auto":
qnode.interface = qml.math.get_interface(*args, *list(kwargs.values()))

jac = qml.jacobian(classical_preprocessing, argnum=wrapper_argnum)(*args, **kwargs)

if qnode.interface == "torch":
import torch

def _jacobian(*args, **kwargs):  # pylint: disable=unused-argument

torch_argnum = (
wrapper_argnum
if wrapper_argnum is not None
else qml.math.get_trainable_indices(args)
)
if np.isscalar(torch_argnum):
jac = jac[torch_argnum]
else:
jac = tuple((jac[idx] for idx in torch_argnum))
return jac

jac = _jacobian(*args, **kwargs)

if qnode.interface in ["jax", "jax-jit"]:
import jax

argnum = 0 if wrapper_argnum is None else wrapper_argnum

def _jacobian(*args, **kwargs):
return jax.jacobian(classical_preprocessing, argnums=argnum)(*args, **kwargs)

jac = _jacobian(*args, **kwargs)

if qnode.interface == "tf":
import tensorflow as tf

def _jacobian(*args, **kwargs):
if np.isscalar(wrapper_argnum):
sub_args = args[wrapper_argnum]
elif wrapper_argnum is None:
sub_args = args
else:
sub_args = tuple((args[i] for i in wrapper_argnum))

gate_params = classical_preprocessing(*args, **kwargs)

jac = tape.jacobian(gate_params, sub_args)
return jac

jac = _jacobian(*args, **kwargs)

if old_interface == "auto":
qnode.interface = "auto"

return jac

return qnode_wrapper


Using PennyLane

Development

API

Internals