stochastic sampler fix

Author: arrjon

bayesflow/experimental/diffusion_model.py

@@ -528,9 +528,9 @@ def velocity(
         self,
         xz: Tensor,
         time: float | Tensor,
+        stochastic_solver: bool,
         conditions: Tensor = None,
         training: bool = False,
-        stochastic_solver: bool = False,
         clip_x: bool = False,
     ) -> Tensor:
         # calculate the current noise level and transform into correct shape
@@ -583,7 +583,7 @@ def _velocity_trace(
         training: bool = False,
     ) -> (Tensor, Tensor):
         def f(x):
-            return self.velocity(x, time=time, conditions=conditions, training=training)
+            return self.velocity(x, time=time, stochastic_solver=False, conditions=conditions, training=training)
 
         v, trace = jacobian_trace(f, xz, max_steps=max_steps, seed=self.seed_generator, return_output=True)
 
@@ -630,7 +630,9 @@ def deltas(time, xz):
             return z, log_density
 
         def deltas(time, xz):
-            return {"xz": self.velocity(xz, time=time, conditions=conditions, training=training)}
+            return {
+                "xz": self.velocity(xz, time=time, stochastic_solver=False, conditions=conditions, training=training)
+            }
 
         state = {"xz": x}
         state = integrate(
@@ -676,12 +678,14 @@ def deltas(time, xz):
 
             return x, log_density
 
-        def deltas(time, xz):
-            return {"xz": self.velocity(xz, time=time, conditions=conditions, training=training)}
-
         state = {"xz": z}
         if integrate_kwargs["method"] == "euler_maruyama":
 
+            def deltas(time, xz):
+                return {
+                    "xz": self.velocity(xz, time=time, stochastic_solver=True, conditions=conditions, training=training)
+                }
+
             def diffusion(time, xz):
                 return {"xz": self.compute_diffusion_term(xz, time=time, training=training)}
 
@@ -692,6 +696,14 @@ def diffusion(time, xz):
                 **integrate_kwargs,
             )
         else:
+
+            def deltas(time, xz):
+                return {
+                    "xz": self.velocity(
+                        xz, time=time, stochastic_solver=False, conditions=conditions, training=training
+                    )
+                }
+
             state = integrate(
                 deltas,
                 state,
@@ -709,6 +721,7 @@ def compute_metrics(
         stage: str = "training",
     ) -> dict[str, Tensor]:
         training = stage == "training"
+        noise_schedule_training_stage = stage == "training" or stage == "validation"
         if not self.built:
             xz_shape = keras.ops.shape(x)
             conditions_shape = None if conditions is None else keras.ops.shape(conditions)
@@ -723,8 +736,10 @@ def compute_metrics(
         # t = keras.ops.cast(i, keras.ops.dtype(x)) / keras.ops.cast(self._timesteps, keras.ops.dtype(x))
 
         # calculate the noise level
-        log_snr_t = expand_right_as(self.noise_schedule.get_log_snr(t, training=training), x)
-        alpha_t, sigma_t = self.noise_schedule.get_alpha_sigma(log_snr_t=log_snr_t, training=training)
+        log_snr_t = expand_right_as(self.noise_schedule.get_log_snr(t, training=noise_schedule_training_stage), x)
+        alpha_t, sigma_t = self.noise_schedule.get_alpha_sigma(
+            log_snr_t=log_snr_t, training=noise_schedule_training_stage
+        )
 
         # generate noise vector
         eps_t = keras.random.normal(ops.shape(x), dtype=ops.dtype(x), seed=self.seed_generator)