code: 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),
              )
  
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 1 for calling a global function with example usage: test_eval_is_deterministic(inp, model_fn, message) without return types