ott.solvers.linear.sinkhorn_lr.LRSinkhorn.init_state