train_step.py

def compute_loss(params,images,labels):
    logits = CNN().apply({'params': params}, images)
    loss = cross_entropy_loss(logits=logits, labels=labels)
    return loss, logits
    
    
@jax.jit
def train_step(state, batch):
  """Train for a single step."""
  images, labels = batch
  (_, logits), grads = jax.value_and_grad(compute_loss, has_aux=True)(state.params,images,labels)
  state = state.apply_gradients(grads=grads)
  metrics = compute_metrics(logits=logits, labels=labels)
  return state, metrics