fix jax all nans

Author: arrjon

bayesflow/utils/integrate.py

@@ -11,6 +11,7 @@
 from bayesflow.types import Tensor
 from bayesflow.utils import filter_kwargs
 from bayesflow.utils.logging import warning
+from keras import backend as K
 
 from . import logging
 
@@ -22,6 +23,15 @@
 STOCHASTIC_METHODS = ["euler_maruyama", "sea", "shark", "two_step_adaptive", "langevin"]
 
 
+def _check_all_nans(state: StateDict):
+    if K.backend() == "jax":
+        return False  # JAX backend does not support checks of the state variables
+    all_nans_flags = []
+    for v in state.values():
+        all_nans_flags.append(keras.ops.all(keras.ops.isnan(v)))
+    return keras.ops.all(keras.ops.stack(all_nans_flags))
+
+
 def euler_step(
     fn: Callable,
     state: StateDict,
@@ -482,11 +492,6 @@ def integrate(
 
 
 ############ SDE Solvers #############
-def _check_all_nans(state: StateDict):
-    all_nans_flags = []
-    for v in state.values():
-        all_nans_flags.append(keras.ops.all(keras.ops.isnan(v)))
-    return keras.ops.all(keras.ops.stack(all_nans_flags))
 
 
 def stochastic_adaptive_step_size_controller(