make optimal transport tests even more reliable

Author: LarsKue

bayesflow/utils/optimal_transport/sinkhorn.py

@@ -16,7 +16,7 @@ def sinkhorn(
     cost: str | Tensor = "euclidean",
     seed: int = None,
     regularization: float = 1.0,
-    max_steps: int = 10000,
+    max_steps: int | None = 10_000,
     tolerance: float = 1e-6,
     numpy: bool = False,
 ) -> (Tensor, Tensor):
@@ -88,7 +88,7 @@ def sinkhorn_indices(
     cost: str | Tensor = "euclidean",
     seed: int = None,
     regularization: float = 1.0,
-    max_steps: int = 1000,
+    max_steps: int | None = 10_000,
     tolerance: float = 1e-6,
     numpy: bool = False,
 ) -> Tensor | np.ndarray:
@@ -111,7 +111,7 @@ def sinkhorn_indices(
         Default: 1.0
 
     :param max_steps: Maximum number of iterations.
-        Default: 1000
+        Default: 10_000
 
     :param tolerance: Absolute tolerance for convergence.
         Default: 1e-6
@@ -164,15 +164,13 @@ def sinkhorn_plan(
 
     :param regularization: Regularization parameter.
         Controls the standard deviation of the Gaussian kernel.
-        Default: 1.0
 
     :param max_steps: Maximum number of iterations.
-        Default: 1000
 
     :param tolerance: Absolute tolerance for convergence.
-        Default: 1e-6
 
     :param numpy: Whether to use numpy or keras backend.
+        Default: False
 
     :return: Tensor of shape (n, m)
         The transport probabilities.

tests/test_utils/test_optimal_transport.py

@@ -9,20 +9,30 @@
 def test_jit_compile():
     import jax
 
-    x = keras.random.normal((100, 100), seed=0)
-    y = keras.random.normal((100, 100), seed=1)
+    x = keras.random.normal((128, 8), seed=0)
+    y = keras.random.normal((128, 8), seed=1)
 
     ot = jax.jit(optimal_transport, static_argnames=["regularization", "seed"])
-    ot(x, y, regularization=0.1, seed=0)
+    ot(x, y, regularization=1.0, seed=0, max_steps=10)
+
+
+def test_shapes():
+    x = keras.random.normal((128, 8), seed=0)
+    y = keras.random.normal((128, 8), seed=1)
+
+    ox, oy = optimal_transport(x, y, regularization=1.0, seed=0, max_steps=10)
+
+    assert keras.ops.shape(ox) == keras.ops.shape(x)
+    assert keras.ops.shape(oy) == keras.ops.shape(y)
 
 
 def test_transport_cost_improves():
-    x = keras.random.normal((100, 2), seed=0)
-    y = keras.random.normal((100, 2), seed=1)
+    x = keras.random.normal((32, 8), seed=0)
+    y = keras.random.normal((32, 8), seed=1)
 
     before_cost = keras.ops.sum(keras.ops.norm(x - y, axis=-1))
 
-    x, y = optimal_transport(x, y, regularization=1e-3, seed=0)
+    x, y = optimal_transport(x, y, regularization=0.1, seed=0, max_steps=None)
 
     after_cost = keras.ops.sum(keras.ops.norm(x - y, axis=-1))
 
@@ -36,7 +46,7 @@ def test_assignment_is_optimal():
     # we could shuffle x and y, but flipping is a more reliable permutation
     y = keras.ops.flip(y, axis=0)
 
-    x, y = optimal_transport(x, y, regularization=1e-6, seed=0)
+    x, y = optimal_transport(x, y, regularization=1e-3, seed=0, max_steps=1000)
 
     cost = keras.ops.sum(keras.ops.norm(x - y, axis=-1))