def make_trainer_state_dict(step,
                            opt_state,
                            history,
                            model_state,
                            input_signature):
  """Creates a trainer state dictionary to save to disk.

  Args:
    step: int, a step number
    opt_state: OptState namedtuple
    history: `trax.history.History`, the history object.
    model_state: A nested structure of the model state.
    input_signature: signature of model inputs.

  Returns:
    A dictionary with the fields of TrainerState and OptState flattened.
  """
  flat_weights, flat_state = tl.flatten_weights_and_state(
      opt_state.weights, model_state)
  return {
      'step': step,
      'flat_weights': flat_weights,
      'slots': opt_state.slots,
      'opt_params': opt_state.opt_params,
      'history': history,
      'flat_state': flat_state,
      'input_signature': input_signature,
      'version_timestamp': 'Jun-18-2020'  # To update in the future if needed.
  }
Creates a trainer state dictionary to save to disk.

  Args:
    step: int, a step number
    opt_state: OptState namedtuple
    history: `trax.history.History`, the history object.
    model_state: A nested structure of the model state.
    input_signature: signature of model inputs.

  Returns:
    A dictionary with the fields of TrainerState and OptState flattened.
  
trax/supervised/trainer_lib.py