Fix optimal transport indexing for tensorflow

Author: LarsKue

bayesflow/networks/flow_matching/flow_matching.py

@@ -82,7 +82,7 @@ def __init__(
             The base probability distribution from which samples are drawn, such as "normal".
             Default is "normal".
         use_optimal_transport : bool, optional
-            Whether to apply optimal transport for improved training stability. Default is False.
+            Whether to apply optimal transport for improved training stability. Default is True.
         loss_fn : str, optional
             The loss function used for training, such as "mse". Default is "mse".
         integrate_kwargs : dict[str, any], optional
@@ -269,7 +269,7 @@ def compute_metrics(
                 )
                 if conditions is not None:
                     # conditions must be resampled along with x1
-                    conditions = conditions[assignments]
+                    conditions = keras.ops.take(conditions, assignments, axis=0)
 
             t = keras.random.uniform((keras.ops.shape(x0)[0],), seed=self.seed_generator)
             t = expand_right_as(t, x0)

bayesflow/utils/optimal_transport/optimal_transport.py

@@ -1,3 +1,5 @@
+import keras
+
 from .log_sinkhorn import log_sinkhorn
 from .sinkhorn import sinkhorn
 
@@ -37,7 +39,7 @@ def optimal_transport(x1, x2, method="log_sinkhorn", return_assignments=False, *
         x1 and x2 in optimal transport permutation order.
     """
     assignments = methods[method.lower()](x1, x2, **kwargs)
-    x2 = x2[assignments]
+    x2 = keras.ops.take(x2, assignments, axis=0)
 
     if return_assignments:
         return x1, x2, assignments