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: Reset the model parameters.
  
  Restores the parameters from the given output_dir if a checkpoint exists,
  otherwise randomly initializes them.
  
  Does not re-jit the model.
  
  Args:
    output_dir: Output directory.
    init_checkpoint: Initial checkpoint (default $output_dir/model.pkl.gz)
code: def reset(self, output_dir, init_checkpoint=None):
      """Reset the model parameters.
  
      Restores the parameters from the given output_dir if a checkpoint exists,
      otherwise randomly initializes them.
  
      Does not re-jit the model.
  
      Args:
        output_dir: Output directory.
        init_checkpoint: Initial checkpoint (default $output_dir/model.pkl.gz)
      """
      self.close()
      self._output_dir = output_dir
      if output_dir is not None:
          tf.io.gfile.makedirs(output_dir)
      else:
          assert not self._should_save_checkpoints
          assert not self._should_write_summaries
  
      # Create summary writers and history.
      if self._should_write_summaries:
          self._train_sw = jaxboard.SummaryWriter(
              os.path.join(output_dir, "train"), enable=self._is_chief
          )
          self._eval_sw = jaxboard.SummaryWriter(
              os.path.join(output_dir, "eval"), enable=self._is_chief
          )
  
      # Reset the train and eval streams.
      self._train_stream = _repeat_stream(
          self._inputs.train_stream, self._n_devices
      )
      # TODO(lukaszkaiser): add an option to evaluate exactly on the full eval
      #   set by adding a padding and stopping the stream when too large.
      self._eval_stream = _repeat_stream(
          self._inputs.eval_stream, self._n_devices
      )
      self._train_eval_stream = _repeat_stream(
          self._inputs.train_eval_stream, self._n_devices
      )
  
      # Restore the training state.
      if output_dir is not None:
          state = load_trainer_state(
              output_dir, self._model_with_loss, init_checkpoint
          )
      else:
          state = TrainerState(
              step=None,
              opt_state=None,
              history=trax_history.History(),
              model_state=None,
          )
      self._step = state.step or 0
      history = state.history
      self._history = history
      if state.opt_state:
          opt_state = state.opt_state
          model_state = state.model_state
      else:
          opt_state, model_state = self._new_opt_state_and_model_state()
          model_state = self._for_n_devices(model_state)
      self._opt_state = OptState(*self._for_n_devices(opt_state))
      self._model_state = model_state
      if not state.opt_state and self._should_save_checkpoints:
          self.save_state(keep=False)
  
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 10 for calling the Trainer obj's underlying member method code with example usage: obj.reset(output_dir, init_checkpoint) without return types