code: 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.
      }
  
imports: import collections
  import functools
  import itertools
  import os
  import sys
  import time
  from absl import logging
  import gin
  import jax
  import tensorflow.compat.v2
  from trax import fastmath
  from trax import jaxboard
  from trax import layers
  from trax import optimizers
  from trax.data import inputs
  from trax.fastmath import numpy
  from trax.fastmath import random
  from trax.layers import base
  from trax.shapes import ShapeDtype
  from trax.supervised import history
  from trax.supervised import lr_schedules
  from trax.supervised import training
label_desc: Use image node_id 9 for calling a global function with example usage: make_trainer_state_dict(step, opt_state, history, model_state, input_signature) and returns: dict