Python

JAX helpers

Optional differentiation around Lattice contraction costs and parameter shifts.

Install zius-quantum[jax]. Helpers wrap parameter-shift expectations so a classical optimizer can sit on top of Lattice or Statefold. This extra does not replace JAX itself.

Parameter shift
from zius.quantum.jaxext import shift_expectation

value, grad = shift_expectation(circuit, parameter_index=0)