qml.jacobian¶

jacobian
(func, argnum=None, method=None, h=None)[source]¶ Returns the Jacobian as a callable function of vectorvalued (functions of) QNodes.
qjit()
and Autograd compatible.Note
When used with
qjit()
, this function currently only supports the Catalyst compiler. Seecatalyst.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 vectorvalued 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 ofqjit()
, this keyword argument has no effect and should not be set. In justintime (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 autodiff. Catalyst supports"parametershift"
and"adjoint"
on internal QNodes. QNodes withdiff_method="finitediff"
are not supported with"auto"
."fd"
represents firstorder finitedifferences for the entire hybrid function.
h (float) – The stepsize value for the finitedifference (
"fd"
) method withinqjit()
decorated functions. This value has no effect in noncompiled 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 builtin 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 atuple
. Thei
th entry of thetuple
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
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 arrayvalued QNode argument with shape
(2, 1, 3)
and outputs the probability of each 2wire basis state, of which there are2**num_wires
= 4. Therefore, the Jacobian of this QNode will be a single array with shape(2, 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:
dev = qml.device("default.qubit", wires=2) @qml.qnode(dev) def circuit(x, y, z): qml.RX(x, wires=0) qml.RY(y, wires=1) qml.RZ(z, wires=0) return tuple(qml.expval(qml.Z(w)) for w in dev.wires) 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 threetuple 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.expval(qml.Z(w)) for w in [0, 1, 2]] 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
andy
, 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 ofx
((2,)
). Similarly, the shape(2, 4)
ofy
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([[1.32116540e07, 1.33781874e07], [4.20735506e01, 4.20735506e01]])
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)
>>> qml.qjit(workflow)(np.array([2.0, 1.0])) array([[0.37120096, 0.45467246], [0.37120096, 0.45467246]])