revert duplicate removal

Author: LeonStadelmann

src/cellflow/solvers/_genot.py

@@ -285,6 +285,7 @@ def _predict_jit(
         x: ArrayLike,
         condition: dict[str, ArrayLike] | None = None,
         rng: ArrayLike | None = None,
+        rng_genot: ArrayLike | None = None,
         **kwargs: Any,
     ) -> ArrayLike | tuple[ArrayLike, diffrax.Solution]:
         kwargs.setdefault("dt0", None)
@@ -295,7 +296,8 @@ def _predict_jit(
         use_mean = rng is None or self.condition_encoder_mode == "deterministic"
         rng = utils.default_prng_key(rng)
         encoder_noise = jnp.zeros(noise_dim) if use_mean else jax.random.normal(rng, noise_dim)
-        latent = self.latent_noise_fn(rng, (x.shape[0],))
+        rng_genot = utils.default_prng_key(rng_genot)
+        latent = self.latent_noise_fn(rng_genot, (x.shape[0],))
 
         def vf(t: float, x: jnp.ndarray, args: tuple[dict[str, jnp.ndarray], jnp.ndarray]) -> jnp.ndarray:
             params = self.vf_state.params