(YN, rng)
| 136 | |
| 137 | # run reverse |
| 138 | def reverse(YN, rng): |
| 139 | Yi = YN |
| 140 | Ybars = [] |
| 141 | with tqdm(range(args.Ndiffuse - 1, 0, -1), desc="Diffusing") as pbar: |
| 142 | for i in pbar: |
| 143 | carry_once = (i, rng, Yi) |
| 144 | (i, rng, Yi), rew = reverse_once(carry_once, None) |
| 145 | Ybars.append(Yi) |
| 146 | # Update the progress bar's suffix to show the current reward |
| 147 | pbar.set_postfix({"rew": f"{rew:.2e}"}) |
| 148 | return jnp.array(Ybars) |
| 149 | |
| 150 | rng_exp, rng = jax.random.split(rng) |
| 151 | Yi = reverse(YN, rng_exp) |
no test coverage detected