(updates)
| 192 | loss_grad = jax.grad(loss, has_aux=True) |
| 193 | |
| 194 | def clip_by_global_norm(updates): |
| 195 | g_norm = optax.global_norm(updates) |
| 196 | trigger = g_norm < max_gradient_norm |
| 197 | return jax.tree_util.tree_map( |
| 198 | lambda t: jnp.where(trigger, t, (t / g_norm) * max_gradient_norm), |
| 199 | updates) |
| 200 | |
| 201 | def training_epoch(training_state: TrainingState, key: PRNGKey, ref_traj: jnp.ndarray, mask: jnp.ndarray): |
| 202 | key, key_grad = jax.random.split(key) |