Using PennyLane

Release news

Development

API

Internals

  1. Docs
  2. qml
  3. qml.jacobian

qml.jacobian

jacobian(func, argnum=None, method=None, h=None)[source]

Returns the Jacobian as a callable function of vector-valued (functions of) QNodes. This function is compatible with Autograd and qjit().

Note

When used with qjit(), this function currently only supports the Catalyst compiler. See catalyst.jacobian() for more details.

Please see the Catalyst quickstart guide, as well as the sharp bits and debugging tips page for an overview of the differences between Catalyst and PennyLane.

Parameters

  • func (function) – A vector-valued Python function or QNode that contains a combination of quantum and classical nodes. The output of the computation must consist of a single NumPy array (if classical) or a tuple of expectation values (if a quantum node)

  • argnum (int or Sequence[int]) – Which argument to take the gradient with respect to. If a sequence is given, the Jacobian corresponding to all marked inputs and all output elements is returned.

  • method (str) –

    Specifies the gradient method when used with the qjit() decorator. Outside of qjit(), this keyword argument has no effect and should not be set. In just-in-time (JIT) mode, this can be any of ["auto", "fd"], where:

    • "auto" represents deferring the quantum differentiation to the method specified by the QNode, while the classical computation is differentiated using traditional auto-diff. Catalyst supports "parameter-shift" and "adjoint" on internal QNodes. QNodes with diff_method="finite-diff" are not supported with "auto".

    • "fd" represents first-order finite-differences for the entire hybrid function.

  • h (float) – The step-size value for the finite-difference ("fd") method within qjit() decorated functions. This value has no effect in non-compiled functions.

Returns

the function that returns the Jacobian of the input function with respect to the arguments in argnum

Return type

function

Note

Due to a limitation in Autograd, this function can only differentiate built-in scalar or NumPy array arguments.

For argnum=None, the trainable arguments are inferred dynamically from the arguments passed to the function. The returned function takes the same arguments as the original function and outputs a tuple. The i-th entry of the tuple has shape (*output shape, *shape of args[argnum[i]]).

If a single trainable argument is inferred, or if a single integer is provided as argnum, the tuple is unpacked and its only entry is returned instead.

Example

Consider the QNode

import pennylane as qml
from pennylane import numpy as np

dev = qml.device("default.qubit", wires=2)

@qml.qnode(dev)
def circuit(weights):
    qml.RX(weights[0, 0, 0], wires=0)
    qml.RY(weights[0, 0, 1], wires=1)
    qml.RZ(weights[1, 0, 2], wires=0)
    return qml.probs()

weights = np.array([[[0.2, 0.9, -1.4]], [[0.5, 0.2, 0.1]]], requires_grad=True)

It has a single array-valued QNode argument with shape (2, 1, 3) and outputs the probability of each 2-wire basis state, of which there are 2**num_wires = 4. Therefore, the Jacobian of this QNode will be a single array with shape (4, 2, 1, 3):

>>> qml.jacobian(circuit)(weights).shape
(4, 2, 1, 3)

On the other hand, consider the following QNode for the same circuit structure:

@qml.qnode(dev)
def circuit(x, y, z):
    qml.RX(x, wires=0)
    qml.RY(y, wires=1)
    qml.RZ(z, wires=0)
    return qml.probs()

x = np.array(0.2, requires_grad=True)
y = np.array(0.9, requires_grad=True)
z = np.array(-1.4, requires_grad=True)

It has three scalar QNode arguments and outputs the probability for each of the 4 basis states. Consequently, its Jacobian will be a three-tuple of arrays with the shape (4,):

>>> jac = qml.jacobian(circuit)(x, y, z)
>>> type(jac)
tuple
>>> for sub_jac in jac:
...     print(sub_jac.shape)
(4,)
(4,)
(4,)

For a more advanced setting of QNode arguments, consider the QNode

dev = qml.device("default.qubit", wires=3)

@qml.qnode(dev)
def circuit(x, y):
    qml.RX(x[0], wires=0)
    qml.RY(y[0, 3], wires=1)
    qml.RX(x[1], wires=2)
    return qml.probs()

x = np.array([0.1, 0.5], requires_grad=True)
y = np.array([[-0.3, 1.2, 0.1, 0.9], [-0.2, -3.1, 0.5, -0.7]], requires_grad=True)

If we do not provide argnum, qml.jacobian will correctly identify both, x and y, as trainable function arguments:

>>> jac = qml.jacobian(circuit)(x, y)
>>> print(type(jac), len(jac))
<class 'tuple'> 2
>>> qml.math.shape(jac[0])
(8, 2)
>>> qml.math.shape(jac[1])
(8, 2, 4)

As we can see, there are two entries in the output, one Jacobian for each QNode argument. The shape (8, 2) of the first Jacobian is the combination of the QNode output shape ((8,)) and the shape of x ((2,)). Similarly, the shape (2, 4) of y leads to a Jacobian shape (8, 2, 4).

Instead, we may choose the output to contain only one of the two entries by providing an iterable as argnum:

>>> jac = qml.jacobian(circuit, argnum=[1])(x, y)
>>> print(type(jac), len(jac))
<class 'tuple'> 1
>>> qml.math.shape(jac)
(1, 8, 2, 4)

Here we included the size of the tuple in the shape analysis, corresponding to the first dimension of size 1.

Finally, we may want to receive the single entry above directly, not as a tuple with a single entry. This is done by providing a single integer as argnum

>>> jac = qml.jacobian(circuit, argnum=1)(x, y)
>>> print(type(jac), len(jac))
<class 'numpy.ndarray'> 8
>>> qml.math.shape(jac)
(8, 2, 4)

As expected, the tuple was unpacked and we directly received the Jacobian of the QNode with respect to y.

We can also compute the Jacobian transformation inside a qjit() decorated program:

dev = qml.device("lightning.qubit", wires=1)

@qml.qjit
def workflow(x):
    @qml.qnode(dev)
    def circuit(x):
        qml.RX(np.pi * x[0], wires=0)
        qml.RY(x[1], wires=0)
        return qml.probs()

    g = qml.jacobian(circuit)
    return g(x)

>>> workflow(np.array([2.0, 1.0]))
Array([[ 3.48786850e-16, -4.20735492e-01],
       [-8.71967125e-17,  4.20735492e-01]], dtype=float64)

You can further compute the Jacobian transformation using other supported differentiation methods by catalyst.jacobian().

@qml.qjit
def workflow(x):
    @qml.qnode(dev)
    def circuit(x):
        qml.RX(np.pi * x[0], wires=0)
        qml.RY(x[1], wires=0)
        return qml.probs()

    g = qml.jacobian(circuit, method="fd", h=0.3)
    return g(x)

>>> workflow(np.array([2.0, 1.0]))
Array([[-0.03996468, -0.42472435],
       [ 0.03996468,  0.42472435]], dtype=float64)
Previous
Next

Downloads

code/api/pennylane.jacobian
 
Download Python script
 
Download Notebook
 
View on GitHub