code: def test_eval_equals_predict_discrete(
      model_fn, vocab_size=10, length=5, batch_size=3
  ):
      """Tests the equivalence of eval and predict modes for discrete models."""
      with fastmath.use_backend(fastmath.Backend.JAX):
          model_slow = model_fn(mode="eval", vocab_size=vocab_size)
          model_fast = model_fn(mode="predict", vocab_size=vocab_size)
          rng = fastmath.random.get_prng(0)
          input_signature = shapes.ShapeDtype((batch_size, 1), np.int32)
          # Given the same rng, both models initialize with the same parameters.
          model_slow.init(input_signature, rng)
          model_fast.init(input_signature, rng)
  
          buf = np.zeros((batch_size, length), dtype=np.int32)
          next_sym = np.zeros((batch_size, 1), dtype=np.int32)
  
          for index in range(length):
              logits_slow = model_slow(buf, rng=rng)
              logits_fast = model_fast(next_sym, rng=rng)
              np.testing.assert_array_almost_equal(
                  logits_slow[:, index, :],
                  logits_fast[:, 0, :],
                  decimal=5,
              )
              next_sym = np.random.randint(
                  vocab_size, size=(batch_size, 1)
              )
              buf[:, index] = next_sym[:, 0]
  
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 4 for calling a global function with example usage: test_eval_equals_predict_discrete(model_fn, vocab_size, length, batch_size) without return types