catalyst.jacobian

jacobian(fn=None, *, method=None, h=None, argnums=None)

A qjit() compatible Jacobian transformation for PennyLane/Catalyst.

This function allows the Jacobian of a hybrid quantum-classical function to be computed within the compiled program. Outside of a compiled function, this function will simply dispatch to its JAX counterpart jax.jacobian. The function f can return any pytree-like shape.

Parameters

• fn (Callable) – a function or a function object to differentiate

• method (str) –

  The method used for differentiation, which 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. Notably, QNodes with diff_method="finite-diff" is 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

• argnums (Tuple[int, List[int]]) – the argument indices to differentiate

Returns

A callable object that computes the Jacobian of the wrapped function for the given

arguments.

Return type

Callable

Raises

ValueError – Invalid method or step size parameters.

Note

Any JAX-compatible optimization library, such as Optax, can be used alongside jacobian for JIT-compatible variational workflows. See the Quick Start for examples.

See also

grad(), value_and_grad().

Example

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

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

    g = jacobian(circuit)
    return g(x)

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

code/api/catalyst.jacobian

 

Download Python script

 

Download Notebook

 

View on GitHub