File size: 1,007 Bytes
0a937d7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 |
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
|