fix ultra-strict convergence criterion in log_sinkhorn_plan

Author: LarsKue

bayesflow/utils/optimal_transport/log_sinkhorn.py

@@ -31,8 +31,8 @@ def contains_nans(plan):
 
     def is_converged(plan):
         # for convergence, the plan should be doubly stochastic
-        conv0 = keras.ops.all(keras.ops.isclose(keras.ops.logsumexp(plan, axis=0), 0.0, rtol=rtol, atol=atol))
-        conv1 = keras.ops.all(keras.ops.isclose(keras.ops.logsumexp(plan, axis=1), 0.0, rtol=rtol, atol=atol))
+        conv0 = keras.ops.all(keras.ops.isclose(keras.ops.logsumexp(plan, axis=0), 0.0, rtol=0.0, atol=atol + rtol))
+        conv1 = keras.ops.all(keras.ops.isclose(keras.ops.logsumexp(plan, axis=1), 0.0, rtol=0.0, atol=atol + rtol))
         return conv0 & conv1
 
     def cond(_, plan):