ott.solvers.linear.sinkhorn.Sinkhorn.kernel_step