remove the is_symbolic_tensor check because this would otherwise skip the whole function for compiled contexts

Author: LarsKue

bayesflow/utils/optimal_transport/log_sinkhorn.py

@@ -1,7 +1,6 @@
 import keras
 
 from .. import logging
-from ..tensor_utils import is_symbolic_tensor
 
 from .euclidean import euclidean
 
@@ -27,9 +26,6 @@ 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))
 

bayesflow/utils/optimal_transport/sinkhorn.py

@@ -3,7 +3,6 @@
 from bayesflow.types import Tensor
 
 from .. import logging
-from ..tensor_utils import is_symbolic_tensor
 
 from .euclidean import euclidean
 
@@ -76,9 +75,6 @@ def sinkhorn_plan(
     # initialize the transport plan from a gaussian kernel
     plan = keras.ops.exp(cost / -(regularization * keras.ops.mean(cost) + 1e-16))
 
-    if is_symbolic_tensor(plan):
-        return plan
-
     def contains_nans(plan):
         return keras.ops.any(keras.ops.isnan(plan))