def test_eval_is_deterministic(inp, model_fn, message=''):
  """Utility method for testing if eval mode is deterministic.

  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.
    message: Optional message to show when outputs of eval/predict mode don't
      match.
  """
  with fastmath.use_backend(fastmath.Backend.JAX):
    model_eval1 = model_fn(mode='eval')
    model_eval2 = model_fn(mode='eval')

    input_signature = shapes.signature(inp)
    model_eval1.init(input_signature)
    model_eval2.init(input_signature)
    model_eval1.save_to_file('/tmp/unique_weights')
    model_eval2.init_from_file('/tmp/unique_weights', weights_only=True,
                               input_signature=input_signature)

    rng = fastmath.random.get_prng(0)
    output_eval1 = model_eval1(inp, rng=rng)
    if not isinstance(output_eval1, (tuple, list)):
      # We will automatically check each and every tensor returned.
      output_eval1 = [output_eval1]

    output_eval2 = model_eval2(inp, rng=rng)
    if not isinstance(output_eval2, (tuple, list)):
      # We will automatically check each and every tensor returned.
      output_eval2 = [output_eval2]

    np.testing.assert_equal(len(output_eval1), len(output_eval2))
    for out1, out2 in zip(output_eval1, output_eval2):
      np.testing.assert_array_almost_equal(
          out1,
          out2,
          decimal=5,
          err_msg='Non-deterministic.{}'.format(message))
Utility method for testing if eval mode is deterministic.

  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.
    message: Optional message to show when outputs of eval/predict mode don't
      match.
  
trax/layers/test_utils.py