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]
Tests the equivalence of eval and predict modes for discrete models.
trax/layers/test_utils.py