value_and_grad¶
- value_and_grad(f: Union[catalyst.jax_tracer.Function, pennylane.workflow.qnode.QNode, Callable, catalyst.jit.QJIT], *, method=None, h=None, argnum=None)[source]¶
A
qjit()
compatible gradient transformation for PennyLane/Catalyst.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
. The functionf
can return any pytree-like shape.Warning
Currently, higher-order differentiation is only supported by the finite-difference method.
- Parameters
f (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 withdiff_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"
) methodargnum (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 JAXopt, can be used alongside
value_and_grad
for JIT-compatible variational workflows. See the Quick Start for examples.See also
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)) g = value_and_grad(circuit) return g(x)
>>> workflow(2.0) (array(0.2), array(-3.14159265))
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.12502201), array(4.34374983))
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), array(4.6))