jcm/losses.py [700:711]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
        data = batch["image"]
        rng = hk.PRNGSequence(rng)

        if isinstance(sde, sde_lib.KVESDE):
            t = random.normal(next(rng), (data.shape[0],)) * 1.2 - 1.2
            t = jnp.exp(t)
        else:
            t = random.uniform(next(rng), (data.shape[0],), minval=eps, maxval=sde.T)

        z = random.normal(next(rng), data.shape)
        mean, std = sde.marginal_prob(data, t)
        perturbed_data = mean + batch_mul(std, z)
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -



jcm/losses.py [813:824]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
        data = batch["image"]
        rng = hk.PRNGSequence(rng)
        # DEBUG: beware of eps!
        if isinstance(sde, sde_lib.KVESDE):
            t = random.normal(next(rng), (data.shape[0],)) * 1.2 - 1.2
            t = jnp.exp(t)
        else:
            t = random.uniform(next(rng), (data.shape[0],), minval=eps, maxval=sde.T)
        # t = random.uniform(next(rng), (data.shape[0],), minval=eps, maxval=sde.T)
        z = random.normal(next(rng), data.shape)
        mean, std = sde.marginal_prob(data, t)
        perturbed_data = mean + batch_mul(std, z)
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -