code: def test_eval_equals_predict(
      inp,
      model_fn,
      seq_axis=1,
      seq_tensor=None,
      init_tokens=3,
      message="",
  ):
      """Utility method for testing equivalence of predict and eval modes.
  
      Args:
        inp: input fed to the model. It can be a tensor, or a tuple of tensors.
        model_fn: function creating a model after calling with `mode` argument.
        seq_axis: axis of sequence_length. In predict mode we iterate over this
          axis. By default `1`, which is 2nd dimension.
        seq_tensor: if `inp` is a tuple, `seq_tensor` is an index of an input tensor
          in this tuple on which we iterate the sequence.
        init_tokens: how many tokens should be passed to the first `predict` call.
        message: Optional message to show when outputs of eval/predict mode don't
          match.
      """
      with fastmath.use_backend(fastmath.Backend.JAX):
          model_eval = model_fn(mode="eval")
          model_predict = model_fn(mode="predict")
  
          input_signature = shapes.signature(inp)
          model_eval.init(input_signature)
          model_predict.init(input_signature)
          model_eval.save_to_file("/tmp/unique_weights")
          model_predict.init_from_file(
              "/tmp/unique_weights",
              weights_only=True,
              input_signature=input_signature,
          )
  
          rng = fastmath.random.get_prng(0)
          output_eval = model_eval(inp, rng=rng)
          if not isinstance(output_eval, (tuple, list)):
              # We will automatically check each and every tensor returned.
              output_eval = [output_eval]
  
          if seq_tensor is None:
              length = inp.shape[seq_axis]
          else:
              length = inp[seq_tensor].shape[seq_axis]
  
          assert (
              length >= init_tokens + 2
          )  # Required to properly test predict mode.
          indices_list = [(0, init_tokens)] + [
              (i, i + 1) for i in range(init_tokens, length)
          ]
  
          for indices in indices_list:
              start, end = indices
              if seq_tensor is None:
                  new_inp = inp.take(
                      indices=np.arange(start, end), axis=seq_axis
                  )
              else:
                  new_inp = list(inp)
                  new_inp[seq_tensor] = new_inp[seq_tensor].take(
                      indices=np.arange(start, end), axis=seq_axis
                  )
  
              output_predict = model_predict(new_inp, rng=rng)
              if not isinstance(output_predict, (tuple, list)):
                  # We will automatically check each and every tensor returned.
                  output_predict = [output_predict]
  
              np.testing.assert_equal(
                  len(output_predict), len(output_eval)
              )
              for outp, oute in zip(output_predict, output_eval):
                  np.testing.assert_array_almost_equal(
                      oute.take(
                          indices=np.arange(start, end), axis=seq_axis
                      ),
                      outp.take(
                          indices=np.arange(0, end - start),
                          axis=seq_axis,
                      ),
                      decimal=5,
                      err_msg="Error on element {} out of {}.{}".format(
                          indices, length, message
                      ),
                  )
  
file_path: trax/trax/layers/test_utils.py
imports: import copy
  import functools
  import numpy
  from trax import fastmath
  from trax import layers
  from trax import shapes
label_desc: Use image node_id 2 for calling a global function with example usage: test_eval_equals_predict(inp, model_fn, seq_axis, seq_tensor, init_tokens, message) without return types