seed in stochastic sampler

Author: arrjon

bayesflow/experimental/diffusion_model.py

@@ -7,7 +7,6 @@
 
 from bayesflow.utils.serialization import serialize, deserialize, serializable
 from bayesflow.types import Tensor, Shape
-import bayesflow as bf
 from bayesflow.networks import InferenceNetwork
 import math
 
@@ -334,7 +333,7 @@ def derivative_log_snr(self, log_snr_t: Tensor, training: bool) -> Tensor:
 
     def get_weights_for_snr(self, log_snr_t: Tensor) -> Tensor:
         """Get weights for the signal-to-noise ratio (snr) for a given log signal-to-noise ratio (lambda)."""
-        return (ops.exp(-log_snr_t) + ops.square(self.sigma_data)) / ops.square(self.sigma_data)
+        return ops.exp(-log_snr_t) + ops.square(self.sigma_data)  # / ops.square(self.sigma_data)
 
     def get_config(self):
         return dict(sigma_data=self.sigma_data, sigma_min=self.sigma_min, sigma_max=self.sigma_max)
@@ -403,7 +402,7 @@ def __init__(
         **kwargs
             Additional keyword arguments passed to the subnet and other components.
         """
-        super().__init__(base_distribution=None, **kwargs)
+        super().__init__(base_distribution="normal", **kwargs)
 
         if isinstance(noise_schedule, str):
             if noise_schedule == "linear":
@@ -433,7 +432,6 @@ def __init__(
         self._clip_max = 5.0
 
         # latent distribution (not configurable)
-        self.base_distribution = bf.distributions.DiagonalNormal()
         self.integrate_kwargs = self.INTEGRATE_DEFAULT_CONFIG | (integrate_kwargs or {})
         self.seed_generator = keras.random.SeedGenerator()
 

bayesflow/utils/integrate.py

@@ -423,18 +423,18 @@ def integrate_stochastic(
     time = start_time
 
     def body(_loop_var, _loop_state):
-        _state, _time = _loop_state
+        _state, _time, _seed = _loop_state
 
         # Generate noise for this step
         _noise = {}
         for key in _state.keys():
             shape = keras.ops.shape(_state[key])
-            _noise[key] = keras.random.normal(shape, seed=seed) * keras.ops.sqrt(keras.ops.abs(step_size))
+            _noise[key] = keras.random.normal(shape, seed=_seed) * keras.ops.sqrt(keras.ops.abs(step_size))
 
         # Perform integration step
         _state, _time, _ = step_fn(_state, _time, step_size, noise=_noise)
 
-        return _state, _time
+        return _state, _time, _seed
 
-    state, time = keras.ops.fori_loop(0, steps, body, (state, time))
+    state, time = keras.ops.fori_loop(0, steps, body, (state, time, seed))
     return state