def train_step(self, batch):
    """Run one training step and update self._opt_state."""
    # Calculate the current optimizer parameters.
    opt_param_updates = self._for_n_devices(
        {'learning_rate': np.array(self.learning_rate)})
    opt_state = self._opt_state
    opt_state.opt_params.update(opt_param_updates)

    # Run the update.
    weights, slots, opt_params = opt_state
    (weights, slots), stat, self._model_state, self._rngs = self._jit_update_fn(
        (weights, slots), self._step, opt_params, batch,
        self._model_state, self._rngs)
    self._opt_state = opt_state._replace(weights=weights, slots=slots)
    if self._should_log_now():
      for name, value in stat.items():
        # TODO(afrozm): value is a scalar, but sometimes JAX is crashing here
        # with a device put array error complaining that it should be an array.
        # On multiple devices, take the mean.
        scalar_value = np.mean(np.array(value))
        self._train_sw.scalar('training/' + name, scalar_value, step=self._step)
    self._step += 1
Run one training step and update self._opt_state.
trax/supervised/trainer_lib.py