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: Evaluate the model and log metrics.
code: def evaluate(self, n_eval_steps):
      """Evaluate the model and log metrics."""
      _, rng = jax_random.split(self._rngs[0])
      # TODO(lukaszkaiser): both model state and parameters by default include
      # the loss layer. Currently, we access the pure-model parameters by just
      # indexing, [0] here. But we should make it more explicit in a better API.
      weights = (self._opt_state.weights[0], self._metrics_weights)
      state = (self._model_state[0], self._metrics_state)
      self.log_step("Evaluation")
      train_eval_slice = itertools.islice(
          self._train_eval_stream, n_eval_steps
      )
      train_metrics, _ = self.evaluation_round(
          train_eval_slice, weights, state, rng
      )
      self.log_metrics(train_metrics, self._train_sw, "train")
      eval_slice = itertools.islice(self._eval_stream, n_eval_steps)
      eval_metrics, _ = self.evaluation_round(
          eval_slice, weights, state, rng
      )
      self.log_metrics(eval_metrics, self._eval_sw, "eval")
      self.log_step("Finished evaluation")
  
      # Save the learning rate in history.
      self._history.append(
          "train",
          "training/learning_rate",
          self._step,
          self.learning_rate,
      )
  
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 13 for calling the Trainer obj's underlying member method code with example usage: obj.evaluate(n_eval_steps) without return types