reintroduce symbolic tensor check in log_sinkhorn

Author: LarsKue

bayesflow/utils/optimal_transport/log_sinkhorn.py

@@ -1,6 +1,7 @@
 import keras
 
 from .. import logging
+from ..tensor_utils import is_symbolic_tensor
 
 from .euclidean import euclidean
 
@@ -26,6 +27,9 @@ def log_sinkhorn_plan(x1, x2, regularization: float = 1.0, rtol=1e-5, atol=1e-8,
 
     log_plan = cost / -(regularization * keras.ops.mean(cost) + 1e-16)
 
+    if is_symbolic_tensor(log_plan):
+        return log_plan
+
     def contains_nans(plan):
         return keras.ops.any(keras.ops.isnan(plan))