jcm/sampling.py [586:598]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
            return samples

        x = jax.lax.fori_loop(0, sde.N - 1, loop_body, x)
        if denoise:
            t = timesteps[sde.N - 1]
            vec_t = jnp.ones((shape[0],)) * t
            denoiser = denoiser_fn(x, vec_t)
            d = 1 / t * x - 1 / t * denoiser
            next_t = timesteps[sde.N]
            samples = x + (next_t - t) * d
        else:
            samples = x
        return samples, sde.N
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -



jcm/sampling.py [627:639]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
            return samples

        x = jax.lax.fori_loop(0, sde.N - 1, loop_body, x)
        if denoise:
            t = timesteps[sde.N - 1]
            vec_t = jnp.ones((shape[0],)) * t
            denoiser = denoiser_fn(x, vec_t)
            d = 1 / t * x - 1 / t * denoiser
            next_t = timesteps[sde.N]
            samples = x + (next_t - t) * d
        else:
            samples = x
        return samples, sde.N
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -