qml.jacobian¶
-
jacobian
(func, argnum=None)[source]¶ Returns the Jacobian as a callable function of vector-valued (functions of) QNodes.
This is a wrapper around the
autograd.jacobian
function.- 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.
- 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 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 tuple(qml.expval(qml.PauliZ(w)) for w in dev.wires) 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 a tuple of two expectation values. Therefore, the Jacobian of this QNode will be a single array with shape(2, 2, 1, 3)
:>>> qml.jacobian(circuit)(weights).shape (2, 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.PauliZ(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 a tuple of two expectation values. Consequently, its Jacobian will be a three-tuple of arrays with the shape
(2,)
:>>> jac = qml.jacobian(circuit)(x, y, z) >>> type(jac) tuple >>> for sub_jac in jac: ... print(sub_jac.shape) (2,) (2,) (2,)
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.PauliZ(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]) (3, 2) >>> qml.math.shape(jac[1]) (3, 2, 4)
As we can see, there are two entries in the output, one Jacobian for each QNode argument. The shape
(3, 2)
of the first Jacobian is the combination of the QNode output shape ((3,)
) and the shape ofx
((2,)
). Similarly, the shape(2, 4)
ofy
leads to a Jacobian shape(3, 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, 3, 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'> 3 >>> qml.math.shape(jac) (3, 2, 4)
As expected, the tuple was unpacked and we directly received the Jacobian of the QNode with respect to
y
.