qml.gradients.jvp

jvp(tape, tangent, gradient_fn, shots=None, gradient_kwargs=None)[source]

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.

Parameters
  • tape (QuantumTape) – quantum tape to differentiate

  • tangent (tensor_like, list) – Gradient-output vector. Must have shape matching the number of trainable parameters.

  • gradient_fn (callable) – the gradient transform to use to differentiate 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 determining the gradients of tapes

Returns

Jacobian vector product. Returns None if the tape has no trainable parameters.

Return type

tensor_like or tuple or None

Example

Consider the following quantum tape with Jax parameters:

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)]
>>> jvp_tapes, fn = qml.gradients.jvp(tape, tangent, qml.gradients.param_shift)

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
(DeviceArray(-0.62073976, dtype=float32), DeviceArray([-0.3259707 ,  0.32597077], dtype=float32))