def get_heun_sampler()

in jcm/sampling.py [0:0]

• 39 lines of code
• 2 McCabe index (conditional complexity)

def get_heun_sampler(sde, model, shape, denoise=True):
    def heun_sampler(rng, state):
        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

            vec_next_t = jnp.ones((shape[0],)) * next_t
            denoiser = denoiser_fn(samples, vec_next_t)
            next_d = 1 / next_t * samples - 1 / next_t * denoiser
            samples = x + (next_t - t) / 2 * (d + next_d)

            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

    return jax.pmap(heun_sampler, axis_name="batch")