def _iterations_implicit_bwd(res, gr):
  """Runs Sinkhorn in backward mode, using implicit differentiation.

  Args:
    res: residual data sent from fwd pass, used for computations below. In this
      case consists in the output itself, as well as inputs against which we
      wish to differentiate.
    gr: gradients w.r.t outputs of fwd pass, here w.r.t size f, g, errors. Note
      that differentiability w.r.t. errors is not handled, and only f, g is
      considered.

  Returns:
    a tuple of gradients: PyTree for geom, one jnp.ndarray for each of a and b.
  """
  f, g, ot_prob, solver = res
  gr = gr[:2]
  return (
      *solver.implicit_diff.gradient(ot_prob, f, g, solver.lse_mode, gr),
      None, None)