code: def epochs(total_steps, steps_to_skip, epoch_steps):
      """Generates the number of steps in each epoch before reaching total_steps.
  
      Args:
        total_steps: int, total number of steps.
        steps_to_skip: int, number of steps to skip because of a restart.
        epoch_steps: iterable of int, numbers of steps in each epoch.
  
      Yields:
        epoch_steps: int, number of steps in this epoch
      """
      steps_to_go = total_steps - steps_to_skip
      epoch_steps = iter(epoch_steps)
  
      # Remove the desired number of steps from the stream.
      for steps_this_epoch in epoch_steps:
          if steps_this_epoch > steps_to_skip:
              # Put back the number of steps left in the unfinished epoch.
              epoch_steps = itertools.chain(
                  [steps_this_epoch - steps_to_skip], epoch_steps
              )
          if steps_this_epoch >= steps_to_skip:
              break
          steps_to_skip -= steps_this_epoch
  
      # Yield the remaining steps per epoch up to total_steps.
      for steps_this_epoch in epoch_steps:
          steps_this_epoch = min(steps_this_epoch, steps_to_go)
          yield steps_this_epoch
          steps_to_go -= steps_this_epoch
          if steps_to_go == 0:
              break
  
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 8 for calling a global function with example usage: epochs(total_steps, steps_to_skip, epoch_steps) without return types