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 potential

  • g (Array) – jnp.ndarray, second potential

  • lse_mode (bool) – bool

Returns:

a jnp.ndarray of size (size of n + m) quantifying deviation to optimality for variables f and g.