def likelihood_fn()

in jcm/likelihood.py [0:0]

• 38 lines of code
• 4 McCabe index (conditional complexity)

    def likelihood_fn(rng, state, data):
        """Compute an unbiased estimate to the log-likelihood in bits/dim.

        Args:
            rng: An array of random states.
            state: Replicated training state for running on multiple devices.
            data: A JAX array of shape [batch size, ...].

        Returns:
            bpd: A JAX array of shape [batch size]. The log-likelihoods on `data` in bits/dim.
            z: A JAX array of the same shape as `data`. The latent representation of `data` under the
                probability flow ODE.
            nfe: An integer. The number of function evaluations used for running the black-box ODE solver.
        """
        div_fn = get_div_fn(lambda x, t: drift_fn(state, x, t))

        rng = hk.PRNGSequence(rng)
        shape = data.shape
        if hutchinson_type == "Gaussian":
            epsilon = jax.random.normal(next(rng), shape)
        elif hutchinson_type == "Rademacher":
            epsilon = jax.random.rademacher(next(rng), shape, dtype=data.dtype)
        else:
            raise NotImplementedError(f"Hutchinson type {hutchinson_type} unknown.")

        ## ODE function for diffrax ODE solver
        def ode_func(t, x, args):
            sample = x[..., :-1]
            vec_t = jnp.ones((sample.shape[0],)) * t
            drift = drift_fn(sample, vec_t)
            logp_grad = div_fn(sample, vec_t, epsilon)
            return jnp.stack([drift, logp_grad], axis=-1)

        term = diffrax.ODETerm(ode_func)
        solver = diffrax.Tsit5()
        stepsize_controller = diffrax.PIDController(rtol=rtol, atol=atol)
        solution = diffrax.diffeqsolve(
            term,
            solver,
            t0=sde.T,
            t1=eps,
            dt0=eps - sde.T,
            y0=jnp.stack([data, jnp.zeros_like((data.shape[0],))], axis=-1),
            stepsize_controller=stepsize_controller,
        )

        nfe = solution.stats["num_steps"]
        z = solution.ys[-1, ..., :-1]
        delta_logp = solution.ys[-1, ..., -1]
        prior_logp = sde.prior_logp(z)
        bpd = -(prior_logp + delta_logp) / np.log(2)
        N = np.prod(shape[1:])
        bpd = bpd / N
        offset = 7.0
        bpd += offset
        return bpd, z, nfe