ott.solvers.linear.implicit_differentiation.ImplicitDiff.gradient

ImplicitDiff.gradient(prob, f, g, lse_mode, gr)

Apply vjp to recover gradient in reverse mode differentiation.

Return type

LinearProblem

Parameters

• prob (ott.problems.linear.linear_problem.LinearProblem) –

• f (jax.Array) –

• g (jax.Array) –

• lse_mode (bool) –

• gr (Tuple[jax.Array, jax.Array]) –