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","."]
code: def __init__(
      self,
      model,
      loss_fn,
      optimizer,
      lr_schedule,
      inputs,
      output_dir=None,
      random_seed=None,
      n_devices=None,
      checkpoints_at=None,
      should_save_checkpoints=True,
      should_write_summaries=True,
      metrics=None,
      checkpoint_highest=None,
      checkpoint_lowest=None,
      init_checkpoint=None,
  ):
      (
          self._is_chief,
          _,
          self._n_devices,
          rng,
      ) = training.init_host_and_devices(n_devices, random_seed)
      self._should_save_checkpoints = (
          should_save_checkpoints and self._is_chief
      )
      self._checkpoints_at = (
          checkpoints_at if checkpoints_at is not None else []
      )
      self._should_write_summaries = should_write_summaries
      if not output_dir:
          self._should_save_checkpoints = False
          self._should_write_summaries = False
      self._checkpoint_highest = checkpoint_highest
      self._checkpoint_lowest = checkpoint_lowest
      self._metrics_dict = (
          metrics if metrics is not None else _DEFAULT_METRICS
      )
      # Inputs is either an Inputs instance or a function that returns it.
      self._inputs = inputs
      if callable(
          inputs
      ):  # If we pass a function, e.g., through gin, call it.
          self._inputs = inputs()
      # Initialize the learning rate to a dummy value. It will be set in reset().
      opt = optimizer(learning_rate=0.0)
  
      # Setup the model.
      model_train = model(mode="train")
      model_predict_eval = model(mode="eval")
      # Should work for fine-tuning of T5.
      if init_checkpoint:
          model_train.init_from_file(init_checkpoint, weights_only=True)
          model_predict_eval.init_from_file(
              init_checkpoint, weights_only=True
          )
      self._model_with_loss = tl.Serial(model_train, loss_fn)
  
      # Setup state.
      rng, init_rng = jax_random.split(rng)
      self._rngs = np.stack(jax_random.split(rng, self._n_devices))
      shapes, dtypes = self._inputs.example_shape_dtype
      input_signature = tuple(
          ShapeDtype(s, d) for (s, d) in zip(shapes, dtypes)
      )
  
      def new_opt_state_and_model_state(rng):
          """Returns optimizer and model states suitable for training a model."""
          weights, state = self._model_with_loss.init(
              input_signature, rng=rng
          )
          (slots, opt_params) = opt.tree_init(weights)
          return (OptState(weights, slots, opt_params), state)
  
      if fastmath.is_backend(fastmath.Backend.JAX):
          # JIT parameter initialization to avoid memory fragmentation
          new_opt_state_and_model_state = fastmath.jit(
              new_opt_state_and_model_state
          )
      self._new_opt_state_and_model_state = (
          lambda: new_opt_state_and_model_state(init_rng)
      )
  
      # Arrange and initialize metrics layers.
      self._metrics = list(sorted(self._metrics_dict.keys()))
      metrics_layers = [self._metrics_dict[m] for m in self._metrics]
      metrics_in_parallel = tl.Branch(*metrics_layers)
      metrics_in_parallel.rng = init_rng
      example_signature = tuple(
          ShapeDtype(s, d)
          for (s, d) in zip(*self._inputs.example_shape_dtype)
      )
      model_predict_eval.init(example_signature)
      self._input_signature = example_signature
      output_signature = model_predict_eval.output_signature(
          example_signature
      )
      m_weights, m_state = metrics_in_parallel.init(output_signature)
      self._metrics_weights = self._for_n_devices(m_weights)
      self._metrics_state = self._for_n_devices(m_state)
  
      # Jit model_predict and update so they're fast.
      self._jit_eval = _jit_predict_fn(
          model_predict_eval, metrics_in_parallel, self._n_devices
      )
      self._jit_update_fn = _jit_update_fn(
          model_train, loss_fn, opt, self._n_devices
      )
  
      self._model_train = model_train
      self._model_predict_eval = model_predict_eval
      self._loss_fn = loss_fn
      self._lr_schedule = lr_schedule
  
      # Those fields will be set in reset().
      self._output_dir = None
      self._train_sw = None
      self._eval_sw = None
      self._history = None
      self._opt_state = None
      self._step = None
      self._model_state = None
      self.reset(output_dir)
  
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 1 to create a new Trainer object  with example: obj = Trainer(model, loss_fn, optimizer, lr_schedule, inputs, output_dir, random_seed, n_devices, checkpoints_at, should_save_checkpoints, should_write_summaries, metrics, checkpoint_highest, checkpoint_lowest, init_checkpoint)