jcm/sde_lib.py [216:231]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
        return mean, std

    def prior_sampling(self, rng, shape):
        return jax.random.normal(rng, shape)

    def prior_logp(self, z):
        shape = z.shape
        N = np.prod(shape[1:])
        logp_fn = lambda z: -N / 2.0 * jnp.log(2 * np.pi) - jnp.sum(z**2) / 2.0
        return jax.vmap(logp_fn)(z)

    def prior_entropy(self, z):
        shape = z.shape
        entropy = jnp.ones(shape) * (0.5 * jnp.log(2 * np.pi) + 0.5)
        entropy = entropy.reshape((z.shape[0], -1))
        return jnp.sum(entropy, axis=-1)
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -



jcm/sde_lib.py [329:344]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
        return mean, std

    def prior_sampling(self, rng, shape):
        return jax.random.normal(rng, shape)

    def prior_logp(self, z):
        shape = z.shape
        N = np.prod(shape[1:])
        logp_fn = lambda z: -N / 2.0 * jnp.log(2 * np.pi) - jnp.sum(z**2) / 2.0
        return jax.vmap(logp_fn)(z)

    def prior_entropy(self, z):
        shape = z.shape
        entropy = jnp.ones(shape) * (0.5 * jnp.log(2 * np.pi) + 0.5)
        entropy = entropy.reshape((z.shape[0], -1))
        return jnp.sum(entropy, axis=-1)
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -