|
|
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 |
|
|
|
