Train the sampler for one iteration.
(dynamics,
x,
optimizer,
loss_fn=l2hmc.compute_loss,
global_step=None)
| 148 | |
| 149 | |
| 150 | def train_one_iter(dynamics, |
| 151 | x, |
| 152 | optimizer, |
| 153 | loss_fn=l2hmc.compute_loss, |
| 154 | global_step=None): |
| 155 | """Train the sampler for one iteration.""" |
| 156 | loss, grads, out, accept_prob = l2hmc.loss_and_grads( |
| 157 | dynamics, x, loss_fn=loss_fn) |
| 158 | optimizer.apply_gradients( |
| 159 | zip(grads, dynamics.trainable_variables), global_step=global_step) |
| 160 | |
| 161 | return loss, out, accept_prob |
| 162 | |
| 163 | |
| 164 | def compute_ac_spectrum(samples_history, target_mean, target_covar): |