class_docstr: Trax trainer.
  
  A trainer allows to make training steps, train for full epochs,
  save the training state and access evaluation data.
class_docstr_tok: ["Trax","trainer",".","A","trainer","allows","to","make","training","steps",",","train","for","full","epochs",",","save","the","training","state","and","access","evaluation","data","."]
docstr: Runs `n_steps` of training, with periodic logging, saving, and evals.
code: def train_epoch(self, n_steps, n_eval_steps):
      """Runs `n_steps` of training, with periodic logging, saving, and evals."""
      # TODO(jonni): Clarify how this method relates to the stricter notion of
      # epoch (training for as many steps as needed for a full pass through the
      # training data).
      print()  # Add visual separator in logs for start of training epoch.
      start_time = time.time()
  
      for _ in range(n_steps):
          batch = next(self._train_stream)
          if (
              self.n_devices > 1
          ):  # TODO(lukaszkaiser): use everywhere if possible.
              batch = _reshape_by_device(batch, self.n_devices)
          self.train_step(batch)
          if self._should_save_now():
              self.save_state(keep=True)
          if self._should_log_now():
              self._train_sw.scalar(
                  "training/learning_rate", self.learning_rate
              )
  
      # At end of n_steps, do bookkeeping, run evals, and save state.
      elapsed_time = time.time() - start_time
      self.log_step(
          "Ran %d train steps in %0.2f secs" % (n_steps, elapsed_time)
      )
      if self._train_sw and n_steps > 1:
          self._train_sw.scalar(
              "training/steps per second",
              n_steps / elapsed_time,
              step=self._step,
          )
          self._train_sw.flush()
      self.evaluate(n_eval_steps)
      if self._eval_sw:
          self._eval_sw.flush()
      if self._should_save_checkpoints:
          self.save_state(keep=False)
      if self._should_save_checkpoints and self._current_step_is_best(
          high=True
      ):
          self.save_state(
              keep=False, prefix="highest_" + self._checkpoint_highest
          )
      if self._should_save_checkpoints and self._current_step_is_best(
          high=False
      ):
          self.save_state(
              keep=False, prefix="lowest_" + self._checkpoint_lowest
          )
  
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 11 for calling the Trainer obj's underlying member method code with example usage: obj.train_epoch(n_steps, n_eval_steps) without return types