enable density computation with adaptive step size solvers

Author: arrjon

bayesflow/networks/diffusion_model/diffusion_model.py

@@ -413,12 +413,6 @@ 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)
@@ -467,12 +461,6 @@ 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,4 +1,3 @@
-import logging
 from collections.abc import Sequence
 
 import keras
@@ -239,12 +238,6 @@ def _forward(
     ) -> 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)
@@ -273,12 +266,6 @@ def _inverse(
     ) -> 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)