ott.solvers.linear.sinkhorn.Sinkhorn.lse_step