def save_state(self, keep, prefix='model'):
    """Save trainer state given a possibly replicated opt_state."""
    opt_state = self._opt_state
    if self.n_devices > 1:
      first_replica = lambda x: x[0]
      opt_state = OptState(*fastmath.nested_map(first_replica, opt_state))
    # This line, while optional, allows JAX to transfer arrays from the device
    # to the host in parallel, which is particularly important for cloud TPU.
    if fastmath.is_backend(fastmath.Backend.JAX):
      opt_state = jax.device_get(opt_state)
    step, history, model_state = self._step, self._history, self._model_state
    output_dir = self._output_dir

    weights_file = os.path.join(output_dir, prefix + '.pkl.gz')

    # This dict will be stored as the model.
    trainer_state_dict = make_trainer_state_dict(
        step, opt_state, history, model_state, self._input_signature)
    self._save_state_dict(trainer_state_dict, weights_file)

    if keep:
      weights_file = os.path.join(output_dir,
                                  '{}_{}.pkl.gz'.format(prefix, step))
      self._save_state_dict(trainer_state_dict, weights_file)
Save trainer state given a possibly replicated opt_state.
trax/supervised/trainer_lib.py