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))
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.
  
trax/layers/test_utils.py