def trainer_state_from_dict(trainer_state_dict, model):
  """Given the trainer state dictionary, returns `TrainerState`."""
  # TODO(afrozm): This becomes simpler if OptState is flattened into
  # TrainerState.
  step = trainer_state_dict['step']
  history = trainer_state_dict['history']
  input_signature = trainer_state_dict['input_signature']
  weights_and_state_sig = model.weights_and_state_signature(input_signature)
  weights, model_state = tl.unflatten_weights_and_state(
      trainer_state_dict['flat_weights'], trainer_state_dict['flat_state'],
      weights_and_state_sig)
  opt_state = OptState(
      weights=weights,
      slots=trainer_state_dict['slots'],
      opt_params=trainer_state_dict['opt_params'])
  return TrainerState(step=step, opt_state=OptState(*opt_state),
                      history=history, model_state=model_state)
Given the trainer state dictionary, returns `TrainerState`.
trax/supervised/trainer_lib.py