| from functools import partial |
|
|
| import jax |
| import jax.numpy as jnp |
|
|
| from .functional import chunk_encode, cache_grad, unchunk_args |
|
|
|
|
| def cache_train_step(loss_fn, state, ss, tt, axis='device'): |
| def encode_with_params(params, **kwargs): |
| return state.apply_fn(params=params, **kwargs) |
|
|
| encode_fn = chunk_encode(partial(encode_with_params, state.params)) |
| grad_fn = cache_grad(encode_with_params) |
|
|
| s_reps = encode_fn(**ss) |
| t_reps = encode_fn(**tt) |
|
|
| @unchunk_args(axis=0, argnums=(0, 1)) |
| def grad_cache_fn(xx, yy): |
| return jnp.mean(loss_fn(xx, yy, axis=axis)) |
| loss, (s_grads, t_grads) = jax.value_and_grad(grad_cache_fn, argnums=(0, 1))(s_reps, t_reps) |
|
|
| grads = jax.tree_map(lambda v: jnp.zeros_like(v), state.params) |
| grads = grad_fn(state.params, grads, s_grads, **ss) |
| grads = grad_fn(state.params, grads, t_grads, **tt) |
|
|
| loss, grads = jax.lax.pmean([loss, grads], axis) |
| new_state = state.apply_gradients(grads=grads) |
| return loss, new_state |
|
|
