ott.solvers.linear.implicit_differentiation.ImplicitDiff.first_order_conditions
ott.solvers.linear.implicit_differentiation.ImplicitDiff.first_order_conditions#
- ImplicitDiff.first_order_conditions(prob, f, g, lse_mode)[source]#
Compute vector of first order conditions for the reg-OT problem.
The output of this vector should be close to zero at optimality. Upon completion of the Sinkhorn forward pass, its norm (using the norm parameter defined using
norm_error
) should be below the threshold parameter.This error will be itself assumed to be close to zero when using implicit differentiation.
- Parameters
prob – definition of the linear optimal transport problem.
f (
Array
) – jnp.ndarray, first potentialg (
Array
) – jnp.ndarray, second potentiallse_mode (
bool
) – bool
- Returns
a jnp.ndarray of size (size of
n + m
) quantifying deviation to optimality for variablesf
andg
.