class_docstr: Evaluates serialized sequence prediction models.
  
  Example: time series prediction. We can serialize a time series into
  a sequence of discrete tokens and model this sequence using an autoregressive
  sequence model, such as Transformer - see
  `trax.rl.serialization_utils.SerializedModel`. Then we can use this callback
  to evaluate long-horizon predictions of such a model.
docstr: Initializes SerializedModelEvaluation.
  
  Args:
    loop: Instance of `trax.supervised.training.Loop` or `None`. Can be set to
      `None` for testing - in such a case, `model` and `eval_task` must be
      provided.
    model: Instance of `trax.rl.serialization_utils.SerializedModel`. Not
      required if `loop` is provided.
    eval_at: When to evaluate. Either int (every how many steps to evaluate),
      or a list of ints (step numbers), or a function int -> bool (step
      predicate).
    eval_task: Instance of `trax.supervised.training.EvalTask` with the
      evaluation data, or None. If not provided, the task will be taken from
      `loop`.
    context_lengths: List of lengths of the context sequence fed into the
      model before starting prediction.
    horizon_lengths: List of lengths of the predicted sequence.
    n_steps: Number of batches to run evaluation for.
    accelerate_model (bool): Whether to wrap the model in `tl.Accelerate`.
docstr_tok: ["Initializes","SerializedModelEvaluation",".","Args",":","loop",":","Instance","of","`","trax.supervised.training.Loop","`","or","`","None","`",".","Can","be","set","to","`","None","`","for","testing","-","in","such","a","case",",","`","model","`","and","`","eval_task","`","must","be","provided",".","model",":","Instance","of","`","trax.rl.serialization_utils.SerializedModel","`",".","Not","required","if","`","loop","`","is","provided",".","eval_at",":","When","to","evaluate",".","Either","int","(","every","how","many","steps","to","evaluate",")",",","or","a","list","of","ints","(","step","numbers",")",",","or","a","function","int","-",">","bool","(","step","predicate",")",".","eval_task",":","Instance","of","`","trax.supervised.training.EvalTask","`","with","the","evaluation","data",",","or","None",".","If","not","provided",",","the","task","will","be","taken","from","`","loop","`",".","context_lengths",":","List","of","lengths","of","the","context","sequence","fed","into","the","model","before","starting","prediction",".","horizon_lengths",":","List","of","lengths","of","the","predicted","sequence",".","n_steps",":","Number","of","batches","to","run","evaluation","for",".","accelerate_model","(","bool",")",":","Whether","to","wrap","the","model","in","`","tl.Accelerate","`","."]
code: def __init__(
      self,
      loop,
      model=None,
      eval_at=1000,
      eval_task=None,
      context_lengths=(1,),
      horizon_lengths=(1,),
      n_steps=1,
      accelerate_model=True,
  ):
      """Initializes SerializedModelEvaluation.
  
      Args:
        loop: Instance of `trax.supervised.training.Loop` or `None`. Can be set to
          `None` for testing - in such a case, `model` and `eval_task` must be
          provided.
        model: Instance of `trax.rl.serialization_utils.SerializedModel`. Not
          required if `loop` is provided.
        eval_at: When to evaluate. Either int (every how many steps to evaluate),
          or a list of ints (step numbers), or a function int -> bool (step
          predicate).
        eval_task: Instance of `trax.supervised.training.EvalTask` with the
          evaluation data, or None. If not provided, the task will be taken from
          `loop`.
        context_lengths: List of lengths of the context sequence fed into the
          model before starting prediction.
        horizon_lengths: List of lengths of the predicted sequence.
        n_steps: Number of batches to run evaluation for.
        accelerate_model (bool): Whether to wrap the model in `tl.Accelerate`.
      """
      super().__init__(loop)
  
      if model is None:
          model = loop.model
  
      observation_serializer = model.observation_serializer
      action_serializer = model.action_serializer
  
      predict_model = model.make_predict_model()
      if accelerate_model:
          predict_model = tl.Accelerate(predict_model)
      self._predict_model = predict_model
      self._obs_serializer = observation_serializer
      self._act_serializer = action_serializer
  
      if isinstance(eval_at, int):
          self._eval_at = lambda step: step % eval_at == 1
      elif hasattr(eval_at, "__in__"):
          self._eval_at = lambda step: step in eval_at
      elif callable(eval_at):
          self._eval_at = eval_at
      else:
          raise TypeError(
              f"Unsupported type for eval_at: {type(eval_at)}."
          )
  
      if eval_task is None:
          if len(loop.eval_tasks) != 1:
              raise ValueError(
                  "If eval_task is not provided, the number of eval_tasks registered "
                  "in Loop must be exactly 1."
              )
          eval_task = loop.eval_tasks[0]
      self._eval_task = eval_task
  
      self._context_lengths = list(sorted(context_lengths))
      self._horizon_lengths = list(sorted(horizon_lengths))
      self._n_steps = n_steps
  
      self._batch_size = eval_task.sample_batch[0].shape[0]
      (_, self._init_state) = predict_model.init(
          shapes.ShapeDtype((self._batch_size, 1), dtype=np.int32)
      )
  
imports: import collections
  import os
  import gin
  import numpy
  from trax import jaxboard
  from trax import layers
  from trax import shapes
  from trax.rl import serialization_utils
  from trax.supervised import decoding
label_desc: Use image node_id 1 to create a new SerializedModelEvaluation object from inherited base classes: TrainingStepCallback with example: obj = SerializedModelEvaluation(loop, model, eval_at, eval_task, context_lengths, horizon_lengths, n_steps, accelerate_model)