fix sinkhorn and log_sinkhorn message formatting for jax by making the warning message worse

Author: LarsKue

bayesflow/utils/optimal_transport/log_sinkhorn.py

@@ -53,22 +53,18 @@ def do_nothing():
         pass
 
     def log_steps():
-        msg = "Log-Sinkhorn-Knopp converged after {:d} steps."
+        msg = "Log-Sinkhorn-Knopp converged after {} steps."
 
         logging.debug(msg, steps)
 
     def warn_convergence():
-        marginals = keras.ops.logsumexp(log_plan, axis=0)
-        deviations = keras.ops.abs(marginals)
-        badness = 100.0 * keras.ops.exp(keras.ops.max(deviations))
+        msg = "Log-Sinkhorn-Knopp did not converge after {} steps."
 
-        msg = "Log-Sinkhorn-Knopp did not converge after {:d} steps (badness: {:.1f}%)."
-
-        logging.warning(msg, max_steps, badness)
+        logging.warning(msg, max_steps)
 
     def warn_nans():
-        msg = "Log-Sinkhorn-Knopp produced NaNs."
-        logging.warning(msg)
+        msg = "Log-Sinkhorn-Knopp produced NaNs after {} steps."
+        logging.warning(msg, steps)
 
     keras.ops.cond(contains_nans(log_plan), warn_nans, do_nothing)
     keras.ops.cond(is_converged(log_plan), log_steps, warn_convergence)

bayesflow/utils/optimal_transport/sinkhorn.py

@@ -102,22 +102,18 @@ def do_nothing():
         pass
 
     def log_steps():
-        msg = "Sinkhorn-Knopp converged after {:d} steps."
+        msg = "Sinkhorn-Knopp converged after {} steps."
 
         logging.info(msg, max_steps)
 
     def warn_convergence():
-        marginals = keras.ops.sum(plan, axis=0)
-        deviations = keras.ops.abs(marginals - 1.0)
-        badness = 100.0 * keras.ops.max(deviations)
+        msg = "Sinkhorn-Knopp did not converge after {}."
 
-        msg = "Sinkhorn-Knopp did not converge after {:d} steps (badness: {:.1f}%)."
-
-        logging.warning(msg, max_steps, badness)
+        logging.warning(msg, max_steps)
 
     def warn_nans():
-        msg = "Sinkhorn-Knopp produced NaNs."
-        logging.warning(msg)
+        msg = "Sinkhorn-Knopp produced NaNs after {} steps."
+        logging.warning(msg, steps)
 
     keras.ops.cond(contains_nans(plan), warn_nans, do_nothing)
     keras.ops.cond(is_converged(plan), log_steps, warn_convergence)