ott.solvers.linear.sinkhorn.Sinkhorn.init_state