fix density computation

Author: arrjon

bayesflow/networks/diffusion_model/diffusion_model.py

@@ -413,6 +413,12 @@ def _forward(
             raise ValueError("Stochastic methods are not supported for forward integration.")
 
         if density:
+            if integrate_kwargs["steps"] == "adaptive":
+                logging.warning(
+                    "Using adaptive integration for density estimation can lead to "
+                    "problems with autodiff. Switching to 200 fixed steps instead."
+                )
+                integrate_kwargs["steps"] = 200
 
             def deltas(time, xz):
                 v, trace = self._velocity_trace(xz, time=time, conditions=conditions, training=training)
@@ -461,6 +467,12 @@ def _inverse(
         if density:
             if integrate_kwargs["method"] in STOCHASTIC_METHODS:
                 raise ValueError("Stochastic methods are not supported for density computation.")
+            if integrate_kwargs["steps"] == "adaptive":
+                logging.warning(
+                    "Using adaptive integration for density estimation can lead to "
+                    "problems with autodiff. Switching to 200 fixed steps instead."
+                )
+                integrate_kwargs["steps"] = 200
 
             def deltas(time, xz):
                 v, trace = self._velocity_trace(xz, time=time, conditions=conditions, training=training)

bayesflow/networks/flow_matching/flow_matching.py

@@ -1,3 +1,4 @@
+import logging
 from collections.abc import Sequence
 
 import keras
@@ -236,14 +237,21 @@ def f(x):
     def _forward(
         self, x: Tensor, conditions: Tensor = None, density: bool = False, training: bool = False, **kwargs
     ) -> Tensor | tuple[Tensor, Tensor]:
+        integrate_kwargs = self.integrate_kwargs | kwargs
         if density:
+            if integrate_kwargs["steps"] == "adaptive":
+                logging.warning(
+                    "Using adaptive integration for density estimation can lead to "
+                    "problems with autodiff. Switching to 200 fixed steps instead."
+                )
+                integrate_kwargs["steps"] = 200
 
             def deltas(time, xz):
                 v, trace = self._velocity_trace(xz, time=time, conditions=conditions, training=training)
                 return {"xz": v, "trace": trace}
 
             state = {"xz": x, "trace": keras.ops.zeros(keras.ops.shape(x)[:-1] + (1,), dtype=keras.ops.dtype(x))}
-            state = integrate(deltas, state, start_time=1.0, stop_time=0.0, **(self.integrate_kwargs | kwargs))
+            state = integrate(deltas, state, start_time=1.0, stop_time=0.0, **integrate_kwargs)
 
             z = state["xz"]
             log_density = self.base_distribution.log_prob(z) + keras.ops.squeeze(state["trace"], axis=-1)
@@ -254,7 +262,7 @@ def deltas(time, xz):
             return {"xz": self.velocity(xz, time=time, conditions=conditions, training=training)}
 
         state = {"xz": x}
-        state = integrate(deltas, state, start_time=1.0, stop_time=0.0, **(self.integrate_kwargs | kwargs))
+        state = integrate(deltas, state, start_time=1.0, stop_time=0.0, **integrate_kwargs)
 
         z = state["xz"]
 
@@ -263,14 +271,21 @@ def deltas(time, xz):
     def _inverse(
         self, z: Tensor, conditions: Tensor = None, density: bool = False, training: bool = False, **kwargs
     ) -> Tensor | tuple[Tensor, Tensor]:
+        integrate_kwargs = self.integrate_kwargs | kwargs
         if density:
+            if integrate_kwargs["steps"] == "adaptive":
+                logging.warning(
+                    "Using adaptive integration for density estimation can lead to "
+                    "problems with autodiff. Switching to 200 fixed steps instead."
+                )
+                integrate_kwargs["steps"] = 200
 
             def deltas(time, xz):
                 v, trace = self._velocity_trace(xz, time=time, conditions=conditions, training=training)
                 return {"xz": v, "trace": trace}
 
             state = {"xz": z, "trace": keras.ops.zeros(keras.ops.shape(z)[:-1] + (1,), dtype=keras.ops.dtype(z))}
-            state = integrate(deltas, state, start_time=0.0, stop_time=1.0, **(self.integrate_kwargs | kwargs))
+            state = integrate(deltas, state, start_time=0.0, stop_time=1.0, **integrate_kwargs)
 
             x = state["xz"]
             log_density = self.base_distribution.log_prob(z) - keras.ops.squeeze(state["trace"], axis=-1)
@@ -281,7 +296,7 @@ def deltas(time, xz):
             return {"xz": self.velocity(xz, time=time, conditions=conditions, training=training)}
 
         state = {"xz": z}
-        state = integrate(deltas, state, start_time=0.0, stop_time=1.0, **(self.integrate_kwargs | kwargs))
+        state = integrate(deltas, state, start_time=0.0, stop_time=1.0, **integrate_kwargs)
 
         x = state["xz"]
 

bayesflow/utils/integrate.py

@@ -22,13 +22,6 @@
 STOCHASTIC_METHODS = ["euler_maruyama", "sea", "shark", "two_step_adaptive", "langevin"]
 
 
-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 euler_step(
     fn: Callable,
     state: StateDict,
@@ -243,22 +236,17 @@ def integrate_fixed(
     step_fn = partial(step_fn, fn, **kwargs, use_adaptive_step_size=False)
     step_size = (stop_time - start_time) / steps
 
-    def cond(_loop_var, _loop_state, _loop_time):
-        all_nans = _check_all_nans(_loop_state)
-        end_now = keras.ops.less(_loop_var, steps)
-        return keras.ops.logical_and(~all_nans, end_now)
-
-    def body(_loop_var, _loop_state, _loop_time):
-        _loop_state, _loop_time, _, _ = step_fn(_loop_state, _loop_time, step_size)
-        return _loop_var + 1, _loop_state, _loop_time
+    def body(_loop_var, _loop_state):
+        _state, _time = _loop_state
+        _state, _time, _, _ = step_fn(_state, _time, step_size)
+        return _state, _time
 
-    _, state, _ = keras.ops.while_loop(
-        cond,
+    state, _ = keras.ops.fori_loop(
+        0,
+        steps,
         body,
-        [0, state, start_time],
+        (state, start_time),
     )
-    if _check_all_nans(state):
-        raise RuntimeError("All values are NaNs in state during integration.")
     return state
 
 
@@ -283,25 +271,18 @@ def integrate_scheduled(
 
     step_fn = partial(step_fn, fn, **kwargs, use_adaptive_step_size=False)
 
-    def cond(_loop_var, _loop_state):
-        all_nans = _check_all_nans(_loop_state)
-        end_now = keras.ops.less(_loop_var, len(steps) - 1)
-        return keras.ops.logical_and(~all_nans, end_now)
-
     def body(_loop_var, _loop_state):
         _time = steps[_loop_var]
         step_size = steps[_loop_var + 1] - steps[_loop_var]
         _loop_state, _, _, _ = step_fn(_loop_state, _time, step_size)
-        return _loop_var + 1, _loop_state
+        return _loop_state
 
-    _, state = keras.ops.while_loop(
-        cond,
+    state = keras.ops.fori_loop(
+        0,
+        keras.ops.shape(steps)[0] - 1,
         body,
-        [0, state],
+        state,
     )
-
-    if _check_all_nans(state):
-        raise RuntimeError("All values are NaNs in state during integration.")
     return state
 
 
@@ -501,6 +482,11 @@ 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(