(carry, unused)
| 192 | |
| 193 | @jax.jit |
| 194 | def reverse_once(carry, unused): |
| 195 | t, rng, Y0 = carry |
| 196 | |
| 197 | # sample from q_i |
| 198 | rng, Y0_rng = jax.random.split(rng) |
| 199 | Y0s = add_noise_batch_to_params(Y0, sigmas[t], Y0_rng) |
| 200 | |
| 201 | # esitimate mu_0tm1 |
| 202 | rng, batch_rng = jax.random.split(rng) |
| 203 | batch_idx = jax.random.choice(batch_rng, Ndata, (Nsample,), replace=False) |
| 204 | train_images_batch = train_images[batch_idx] |
| 205 | train_labels_batch = train_labels[batch_idx] |
| 206 | l = jax.vmap(loss, in_axes=(0, None))(Y0s, (train_images_batch, train_labels_batch)) |
| 207 | Js = -l |
| 208 | # Js = -jax.vmap(eval_fn)(Y0s) |
| 209 | logp0 = (Js - Js.mean()) / Js.std() / temp_sample |
| 210 | weights = jax.nn.softmax(logp0) |
| 211 | Y0 = get_params_batch_weighted_sum(Y0s, weights) |
| 212 | |
| 213 | return (t - 1, rng, Y0), Js.mean() |
| 214 | |
| 215 | |
| 216 | Y0 = params |
no test coverage detected