catalyst.value_and_grad

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

A qjit()-compatible transformation for returning the result and gradient of a function.

This function allows the value and the gradient 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.value_and_grad.

Note that value_and_grad can be more efficient, and reduce overall quantum executions, compared to separately executing the function and then computing its gradient.

Warning

Currently, higher-order differentiation is only supported by the finite-difference method.

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 value and gradient of the wrapped function for the given arguments.

Return type

Callable

Raises

• ValueError – Invalid method or step size parameters.

• DifferentiableCompilerError – Called on a function that doesn’t return a single scalar.

Note

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

See also

grad(), jacobian()

Example 1 (Classical preprocessing)

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

@qjit
def workflow(x):
    @qml.qnode(dev)
    def circuit(x):
        qml.RX(jnp.pi * x, wires=0)
        return qml.expval(qml.PauliY(0))
    return value_and_grad(circuit)(x)

>>> workflow(0.2)
(Array(-0.58778525, dtype=float64),
(Array(-0.58778525, dtype=float64), Array(-2.54160185, dtype=float64)))

Example 2 (Classical preprocessing and postprocessing)

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

@qjit
def value_and_grad_loss(theta):
    @qml.qnode(dev, diff_method="adjoint")
    def circuit(theta):
        qml.RX(jnp.exp(theta ** 2) / jnp.cos(theta / 4), wires=0)
        return qml.expval(qml.PauliZ(wires=0))

    def loss(theta):
        return jnp.pi / jnp.tanh(circuit(theta))

    return catalyst.value_and_grad(loss, method="auto")(theta)

>>> value_and_grad_loss(1.0)
(Array(-4.2622289, dtype=float64), Array(5.04324559, dtype=float64))

Example 3 (Purely classical functions)

def square(x: float):
    return x ** 2

@qjit
def dsquare(x: float):
    return catalyst.value_and_grad(square)(x)

>>> dsquare(2.3)
(Array(5.29, dtype=float64), Array(4.6, dtype=float64))

code/api/catalyst.value_and_grad

 

Download Python script

 

Download Notebook

 

View on GitHub