geope.jax.jacobian

geope.jax.jacobian.Ui(x, basis)

Compute a unitary from a linear combination of Hermitian basis matrices.

Parameters:

Name	Type	Description	Default
x	Array	Coefficient vector of shape (K,).	required
basis	Array	Array of Hermitian matrices of shape (K, d, d).	required

Returns:

Type	Description
Array	A unitary matrix of shape (d, d).

geope.jax.jacobian.get_Ui_fn(basis)

Create a partial unitary function with a fixed basis.

Parameters:

Name	Type	Description	Default
basis	Array	Array of Hermitian matrices of shape (K, d, d).	required

Returns:

Type	Description
Callable[[Array], Array]	A callable that takes a coefficient vector and returns
Callable[[Array], Array]	the corresponding unitary.

geope.jax.jacobian.manual_jacobian(params, Ui_fn, jac_fn)

Compute the full Jacobian of the product unitary manually.

The product unitary follows the convention of :func:geope.engine.compute_matrices_params_list_fn, where each gate is left-multiplied onto the accumulator,

\[U = U_{G-1} \cdots U_1 U_0, \qquad U_i = \exp\!\Big(i \sum_k x_{i,k} G_k\Big).\]

The derivative with respect to a parameter of gate \(i\) leaves every other gate untouched, so it is a product with a single factor replaced by the per-gate derivative:

\[\frac{\partial U}{\partial x_{i,k}} = \underbrace{U_{G-1} \cdots U_{i+1}}_{L_i}\, \frac{\partial U_i}{\partial x_{i,k}}\, \underbrace{U_{i-1} \cdots U_0}_{R_i}.\]

Both the left (\(L_i\), exclusive suffix product) and right (\(R_i\), exclusive prefix product) partial products are obtained in \(O(G)\) matrix multiplications with two jax.lax.scan passes, after which the per-gate derivative blocks are combined with a single vectorised einsum. This is the equivalent of differentiating the whole sequence with autodiff, but built explicitly from the per-gate derivative jac_fn.

Parameters:

Name	Type	Description	Default
params	Array	Parameter Array of shape (G, K).	required
Ui_fn	Callable[[Array], Array]	Callable mapping a coefficient Array to a unitary Array.	required
jac_fn	Callable[[Array], Array]	Callable computing the per-gate Jacobian Array of shape (d, d, K) (e.g. :func:geope.jax.dexpm).	required

Returns:

Type	Description
Array	An Array of shape (G, d, d, K) containing the full Jacobian.

geope.jax.jacobian.get_jacobian_propagator(gate_basis, hermitian=True)

Create a JIT-compiled manual Jacobian function for a given gate basis.

The per-gate derivative uses the spectral method (geope.jax.dexpm_eig), and the returned function is wrapped in jax.jit so it is compiled once and reused across calls (rather than retracing on every invocation).

Parameters:

Name	Type	Description	Default
gate_basis	Array	Array of Hermitian basis matrices of shape (K, d, d).	required
hermitian	bool	Assume real parameters (skew-Hermitian generators) and use the faster eigh-based per-gate derivative. Set False for complex-valued parameters.	True

Returns:

Type	Description
Callable[[Array], Array]	A Callable[[Array], Array] that accepts a parameter array
Callable[[Array], Array]	of shape (G, K) and returns the Jacobian of shape
Callable[[Array], Array]	(G, d, d, K).