add comment about convergence check

Author: LarsKue

bayesflow/utils/optimal_transport/log_sinkhorn.py

@@ -31,6 +31,8 @@ def contains_nans(plan):
 
     def is_converged(plan):
         # for convergence, the plan should be doubly stochastic
+        # NOTE: for small atol and rtol, using rtol_log=0.0 and atol_log=atol + rtol
+        # is equivalent to the convergence check in the unstabilized version
         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