(seed=0)
| 81 | |
| 82 | |
| 83 | def run_exp(seed=0): |
| 84 | rng = jax.random.PRNGKey(seed) |
| 85 | mu_0t = jnp.zeros([Nsample, dim]) + 1.0 * jax.random.normal(rng, (Nsample, dim)) |
| 86 | _, _ = reverse_once((0, rng, mu_0t), None) # to compile |
| 87 | xs, ys = [], [] |
| 88 | with tqdm(range(Ndiffuse - 1, 0, -1), desc="Diffusing") as pbar: |
| 89 | for t in pbar: |
| 90 | carry_once = (t, rng, mu_0t) |
| 91 | (t, rng, mu_0t), J = reverse_once(carry_once, None) |
| 92 | xs.append((Ndiffuse - 1 - t) * Nsample) |
| 93 | ys.append(J) |
| 94 | pbar.set_postfix({"rew": f"{J:.2e}"}) |
| 95 | return jnp.array(xs), jnp.array(ys) |
| 96 | |
| 97 | |
| 98 | if __name__ == "__main__": |
no test coverage detected