jcm/sampling.py [558:579]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
        denoiser_fn = mutils.get_denoiser_fn(
            sde, model, state.params_ema, state.model_state, train=False
        )

        rng = hk.PRNGSequence(rng)
        x = sde.prior_sampling(next(rng), shape)
        timesteps = (
            sde.t_max ** (1 / sde.rho)
            + jnp.arange(sde.N)
            / (sde.N - 1)
            * (sde.t_min ** (1 / sde.rho) - sde.t_max ** (1 / sde.rho))
        ) ** sde.rho
        timesteps = jnp.concatenate([timesteps, jnp.array([0.0])])

        def loop_body(i, val):
            x = val
            t = timesteps[i]
            vec_t = jnp.ones((shape[0],)) * t
            denoiser = denoiser_fn(x, vec_t)
            d = 1 / t * x - 1 / t * denoiser
            next_t = timesteps[i + 1]
            samples = x + (next_t - t) * d
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -



jcm/sampling.py [605:626]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
        denoiser_fn = mutils.get_denoiser_fn(
            sde, model, state.params_ema, state.model_state, train=False
        )

        rng = hk.PRNGSequence(rng)
        x = sde.prior_sampling(next(rng), shape)
        timesteps = (
            sde.t_max ** (1 / sde.rho)
            + jnp.arange(sde.N)
            / (sde.N - 1)
            * (sde.t_min ** (1 / sde.rho) - sde.t_max ** (1 / sde.rho))
        ) ** sde.rho
        timesteps = jnp.concatenate([timesteps, jnp.array([0.0])])

        def loop_body(i, val):
            x = val
            t = timesteps[i]
            vec_t = jnp.ones((shape[0],)) * t
            denoiser = denoiser_fn(x, vec_t)
            d = 1 / t * x - 1 / t * denoiser
            next_t = timesteps[i + 1]
            samples = x + (next_t - t) * d
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -