Repository: keras-team/keras-nlp

Issue:
Make Decoding Functions Graph-compatible
Made an attempt to do the above here: https://colab.research.google.com/drive/1PBMzeBd-HyFE0o4VXwk19-kqHIhOZM49?usp=sharing. Ran into an issue: 

```
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-52-cee4282e867d> in <module>()
      1 inputs = keras.Input(shape=(), dtype="string")
----> 2 translated = decode_sequences(inputs)
      3 translated = tf.strings.unicode_decode(input=translated, input_encoding="UTF-8")
      4 translated = tf.strings.unicode_encode(input=translated, output_encoding="UTF-8")
      5 translated = tf.strings.regex_replace(translated, "\[PAD\]", "")

3 frames
<ipython-input-48-ace559a1a08f> in decode_sequences(input_sentences)
    151         prompt,
    152         max_length=40,
--> 153         end_token_id=spa_tokenizer.token_to_id("[END]"),
    154     )
    155     generated_sentences = spa_tokenizer.detokenize(generated_tokens)

<ipython-input-48-ace559a1a08f> in greedy_search(token_probability_fn, prompt, max_length, end_token_id, pad_token_id)
    123     if end_token_id is not None:
    124         prompt = mask_tokens_after_end_token(
--> 125             prompt, max_length, end_token_id, pad_token_id
    126         )
    127 

<ipython-input-48-ace559a1a08f> in mask_tokens_after_end_token(prompt, max_length, end_token_id, pad_token_id)
     36     # Build a mask including end_token and replace tokens after end_token
     37     # with `pad_token_id`.
---> 38     valid_indices = tf.sequence_mask(lengths=end_indices + 1, maxlen=max_length)
     39     return tf.where(valid_indices, prompt, pad_token_id)
     40 

/usr/local/lib/python3.7/dist-packages/tensorflow/python/util/traceback_utils.py in error_handler(*args, **kwargs)
    151     except Exception as e:
    152       filtered_tb = _process_traceback_frames(e.__traceback__)
--> 153       raise e.with_traceback(filtered_tb) from None
    154     finally:
    155       del filtered_tb

<__array_function__ internals> in result_type(*args, **kwargs)

TypeError: Cannot interpret '<KerasTensor: shape=(None,) dtype=int64 (created by layer 'tf.where_11')>' as a data type
```

This is most probably a bug in Keras. This is the origin of the error: https://github.com/tensorflow/tensorflow/blob/v2.9.1/tensorflow/python/ops/array_ops.py#L4510. All these `.dtype` calls throw an error. I'll open an issue in the Keras repository!

Also, would be glad to work on making the functions graph compatible :)


Discussion:
Thanks for opening! Definitely this is a good thing to work on, getting efficient function tracing support for sequence generation is important. There's actually two variants to this requirement.
 - We should be able to trace functions normally (e.g. in a keras model, or with @tf.function).
 - We should be able to trace functions with XLA (e.g. a keras model or @tf.function with jit_compile=True).

IIUC jit compilation will come with a few more restrictions on the graph code. We can do this incrementally too, start with jit_compile, add it later.

Resolving patch:
diff --git a/keras_nlp/utils/text_generation.py b/keras_nlp/utils/text_generation.py
--- a/keras_nlp/utils/text_generation.py
+++ b/keras_nlp/utils/text_generation.py
@@ -18,6 +18,8 @@
 from absl import logging
 from tensorflow import keras
 
+# TODO (@chenmoneygithub): Refactor code to reuse snippets.
+
 
 def validate_prompt(prompt):
     """Helper function to validate input to text_generation utils."""
@@ -52,7 +54,11 @@ def mask_tokens_after_end_token(prompt, max_length, end_token_id, pad_token_id):
     # Find index of first end_token_id.
     end_indices = tf.math.argmax(prompt == end_token_id, -1)
     # Use max_length if no `end_token_id` is found.
-    end_indices = tf.where(end_indices == 0, max_length, end_indices)
+    end_indices = tf.where(
+        end_indices == 0,
+        tf.cast(max_length, dtype=end_indices.dtype),
+        end_indices,
+    )
     # Build a mask including end_token and replace tokens after end_token
     # with `pad_token_id`.
     valid_indices = tf.sequence_mask(end_indices + 1, maxlen=max_length)
@@ -128,14 +134,6 @@ def token_probability_fn(inputs):
     ```
 
     """
-    if not tf.executing_eagerly():
-        raise RuntimeError(
-            "`keras_nlp.utils.greedy_search` currently requires an eager "
-            "execution context. Please call `greedy_search` outside "
-            "tf.function or run `tf.config.run_functions_eagerly(True)` to run "
-            "tf.function in eager mode."
-        )
-
     prompt = validate_prompt(prompt)
 
     input_is_1d = prompt.shape.rank == 1
@@ -143,14 +141,41 @@ def token_probability_fn(inputs):
         prompt = prompt[tf.newaxis, :]
     validate_token_probability_fn(token_probability_fn, prompt)
 
-    i = prompt.shape[1]
-    while i < max_length:
-        # If the prompt has reached our desired length, exit while loop.
-        pred = token_probability_fn(prompt)
+    shape = tf.shape(prompt)
+    batch_size = shape[0]
+    length = shape[1]
+
+    # Pad the prompt with `pad_token_id` to `max_length`.
+    padding = tf.fill((batch_size, max_length - length), pad_token_id)
+    prompt = tf.concat((prompt, padding), axis=1)
+
+    def one_step(length, prompt):
+        pred = token_probability_fn(prompt[:, :length])
         next_token = tf.cast(tf.argmax(pred, axis=-1), dtype=prompt.dtype)
+
         # Append the next token to current sequence.
-        prompt = tf.concat([prompt, next_token[:, tf.newaxis]], axis=-1)
-        i += 1
+        def add_token(args):
+            sequence, token = args
+            return tf.tensor_scatter_nd_update(
+                tensor=sequence, indices=[[length]], updates=[token]
+            )
+
+        prompt = tf.map_fn(
+            fn=add_token,
+            elems=(prompt, next_token),
+            fn_output_signature=tf.TensorSpec(
+                shape=(max_length), dtype=prompt.dtype
+            ),
+        )
+        length += 1
+        return (length, prompt)
+
+    # Run a while loop till text of length `max_length` has been generated.
+    length, prompt = tf.while_loop(
+        cond=lambda length, _: tf.less(length, max_length),
+        body=one_step,
+        loop_vars=(length, prompt),
+    )
 
     if end_token_id is not None:
         prompt = mask_tokens_after_end_token(
@@ -382,33 +407,55 @@ def token_probability_fn(inputs):
     ```
 
     """
-    if not tf.executing_eagerly():
-        raise RuntimeError(
-            "`keras_nlp.utils.random_sampling` currently requires an eager "
-            "execution context. Please call `random_sampling` outside "
-            "tf.function or run `tf.config.run_functions_eagerly(True)` to run "
-            "tf.function in eager mode."
-        )
-
     prompt = validate_prompt(prompt)
     input_is_1d = prompt.shape.rank == 1
     if input_is_1d:
         prompt = prompt[tf.newaxis, :]
     validate_token_probability_fn(token_probability_fn, prompt)
 
-    i = prompt.shape[1]
-    while i < max_length:
-        # If the prompt has reached our desired length, exit while loop.
+    shape = tf.shape(prompt)
+    batch_size = shape[0]
+    length = shape[1]
+
+    # Pad the prompt with `pad_token_id` to `max_length`.
+    padding = tf.fill((batch_size, max_length - length), pad_token_id)
+    prompt = tf.concat((prompt, padding), axis=1)
+
+    def one_step(length, prompt):
         pred = token_probability_fn(prompt)
         if from_logits:
             pred = keras.activations.softmax(pred, axis=-1)
-        next_token = tf.cast(
-            tf.random.categorical(tf.math.log(pred), 1, seed=seed),
-            dtype=prompt.dtype,
+        next_token = tf.squeeze(
+            tf.cast(
+                tf.random.categorical(tf.math.log(pred), 1, seed=seed),
+                dtype=prompt.dtype,
+            ),
+            axis=1,
         )
+
         # Append the next token to current sequence.
-        prompt = tf.concat([prompt, next_token], axis=-1)
-        i += 1
+        def add_token(args):
+            sequence, token = args
+            return tf.tensor_scatter_nd_update(
+                tensor=sequence, indices=[[length]], updates=[token]
+            )
+
+        prompt = tf.map_fn(
+            fn=add_token,
+            elems=(prompt, next_token),
+            fn_output_signature=tf.TensorSpec(
+                shape=(max_length), dtype=prompt.dtype
+            ),
+        )
+        length += 1
+        return (length, prompt)
+
+    # Run a while loop till text of length `max_length` has been generated.
+    length, prompt = tf.while_loop(
+        cond=lambda length, _: tf.less(length, max_length),
+        body=one_step,
+        loop_vars=(length, prompt),
+    )
 
     if end_token_id is not None:
         prompt = mask_tokens_after_end_token(
@@ -497,14 +544,6 @@ def token_probability_fn(inputs):
     ```
 
     """
-    if not tf.executing_eagerly():
-        raise RuntimeError(
-            "`keras_nlp.utils.top_k_search` currently requires an eager "
-            "execution context. Please call `top_k_search` outside "
-            "tf.function or run `tf.config.run_functions_eagerly(True)` to run "
-            "tf.function in eager mode."
-        )
-
     prompt = validate_prompt(prompt)
     input_is_1d = prompt.shape.rank == 1
     if input_is_1d:
@@ -522,8 +561,15 @@ def token_probability_fn(inputs):
         )
         k = pred.shape[1]
 
-    i = prompt.shape[1]
-    while i < max_length:
+    shape = tf.shape(prompt)
+    batch_size = shape[0]
+    length = shape[1]
+
+    # Pad the prompt with `pad_token_id` to `max_length`.
+    padding = tf.fill((batch_size, max_length - length), pad_token_id)
+    prompt = tf.concat((prompt, padding), axis=1)
+
+    def one_step(length, prompt):
         pred = token_probability_fn(prompt)
         if from_logits:
             pred = keras.activations.softmax(pred, axis=-1)
@@ -534,12 +580,34 @@ def token_probability_fn(inputs):
         next_token = tf.random.categorical(
             tf.math.log(top_k_pred), 1, seed=seed
         )
+
         # Rearrange to get the next token idx from the original order.
         next_token = tf.gather_nd(top_k_indices, next_token, batch_dims=1)
         next_token = tf.cast(next_token, dtype=prompt.dtype)
+
         # Append the next token to current sequence.
-        prompt = tf.concat([prompt, next_token[:, tf.newaxis]], axis=-1)
-        i += 1
+        def add_token(args):
+            sequence, token = args
+            return tf.tensor_scatter_nd_update(
+                tensor=sequence, indices=[[length]], updates=[token]
+            )
+
+        prompt = tf.map_fn(
+            fn=add_token,
+            elems=(prompt, next_token),
+            fn_output_signature=tf.TensorSpec(
+                shape=(max_length), dtype=prompt.dtype
+            ),
+        )
+        length += 1
+        return (length, prompt)
+
+    # Run a while loop till text of length `max_length` has been generated.
+    length, prompt = tf.while_loop(
+        cond=lambda length, _: tf.less(length, max_length),
+        body=one_step,
+        loop_vars=(length, prompt),
+    )
 
     if end_token_id is not None:
         prompt = mask_tokens_after_end_token(
@@ -630,13 +698,6 @@ def token_probability_fn(inputs):
     ```
 
     """
-    if not tf.executing_eagerly():
-        raise RuntimeError(
-            "`keras_nlp.utils.top_p_search` currently requires an eager "
-            "execution context. Please call `top_p_search` outside "
-            "tf.function or run `tf.config.run_functions_eagerly(True)` to run "
-            "tf.function in eager mode."
-        )
     if p <= 0 or p >= 1:
         raise ValueError(
             f"`p` should be in the range (0, 1). Received: `p={p}`."
@@ -648,9 +709,15 @@ def token_probability_fn(inputs):
         prompt = prompt[tf.newaxis, :]
     validate_token_probability_fn(token_probability_fn, prompt)
 
-    i = prompt.shape[1]
-    while i < max_length:
-        # If the prompt has reached our desired length, exit while loop.
+    shape = tf.shape(prompt)
+    batch_size = shape[0]
+    length = shape[1]
+
+    # Pad the prompt with `pad_token_id` to `max_length`.
+    padding = tf.fill((batch_size, max_length - length), pad_token_id)
+    prompt = tf.concat((prompt, padding), axis=1)
+
+    def one_step(length, prompt):
         pred = token_probability_fn(prompt)
         if from_logits:
             pred = keras.activations.softmax(pred, axis=-1)
@@ -679,9 +746,30 @@ def token_probability_fn(inputs):
             sorted_indices, sorted_next_token, batch_dims=1
         )
         next_token = tf.cast(next_token, dtype=prompt.dtype)
+
         # Append the next token to current sequence.
-        prompt = tf.concat([prompt, next_token[:, tf.newaxis]], axis=-1)
-        i += 1
+        def add_token(args):
+            sequence, token = args
+            return tf.tensor_scatter_nd_update(
+                tensor=sequence, indices=[[length]], updates=[token]
+            )
+
+        prompt = tf.map_fn(
+            fn=add_token,
+            elems=(prompt, next_token),
+            fn_output_signature=tf.TensorSpec(
+                shape=(max_length), dtype=prompt.dtype
+            ),
+        )
+        length += 1
+        return (length, prompt)
+
+    # Run a while loop till text of length `max_length` has been generated.
+    length, prompt = tf.while_loop(
+        cond=lambda length, _: tf.less(length, max_length),
+        body=one_step,
+        loop_vars=(length, prompt),
+    )
 
     if end_token_id is not None:
         prompt = mask_tokens_after_end_token(


Test patch:
diff --git a/keras_nlp/utils/text_generation_test.py b/keras_nlp/utils/text_generation_test.py
--- a/keras_nlp/utils/text_generation_test.py
+++ b/keras_nlp/utils/text_generation_test.py
@@ -17,6 +17,7 @@
 
 import numpy as np
 import tensorflow as tf
+from absl.testing import parameterized
 from tensorflow import keras
 
 from keras_nlp.utils.text_generation import beam_search
@@ -26,7 +27,7 @@
 from keras_nlp.utils.text_generation import top_p_search
 
 
-class GreedySearchTextGenerationTest(tf.test.TestCase):
+class GreedySearchTextGenerationTest(tf.test.TestCase, parameterized.TestCase):
     def setUp(self):
         super().setUp()
         self.vocab_size = 10
@@ -81,7 +82,7 @@ def test_generate_with_ragged_prompt(self):
     def test_assert_generation_is_correct(self):
         def token_probability_fn(inputs):
             batch_size = inputs.shape[0]
-            prob = tf.constant([[0.01, 0.01, 0.08, 0.9]])
+            prob = tf.constant([[0.0, 0.0, 0.0, 1.0]])
             return tf.repeat(prob, batch_size, axis=0)
 
         batch_size = 10
@@ -97,7 +98,7 @@ def token_probability_fn(inputs):
     def test_end_token_id(self):
         def token_probability_fn(inputs):
             batch_size = inputs.shape[0]
-            prob = tf.constant([[0.01, 0.01, 0.08, 0.9]])
+            prob = tf.constant([[0.0, 0.0, 0.0, 1.0]])
             return tf.repeat(prob, batch_size, axis=0)
 
         max_length = 5
@@ -113,6 +114,69 @@ def token_probability_fn(inputs):
         expected_outputs = tf.concat([inputs, expected_outputs], axis=1)
         self.assertAllEqual(outputs, expected_outputs)
 
+    @parameterized.named_parameters(
+        ("jit_compile_false", False), ("jit_compile_true", True)
+    )
+    def test_model_compile(self, jit_compile):
+        def token_probability_fn(inputs):
+            prob = tf.constant([[0.0, 0.0, 0.0, 1.0]])
+            return tf.repeat(prob, 2, axis=0)
+
+        max_length = 5
+
+        class TestModel(tf.keras.Model):
+            def call(self, inputs):
+                generated = greedy_search(
+                    token_probability_fn,
+                    inputs,
+                    max_length=max_length,
+                    end_token_id=2,
+                    pad_token_id=0,
+                )
+                return generated
+
+        inputs = tf.constant([[0, 1], [1, 2]])
+        expected_outputs = tf.tile([[3], [0]], [1, max_length - 2])
+        expected_outputs = tf.concat([inputs, expected_outputs], axis=1)
+
+        model = TestModel()
+        model.compile(jit_compile=jit_compile)
+
+        outputs = model.predict(inputs)
+        self.assertAllEqual(outputs, expected_outputs)
+
+    @parameterized.named_parameters(
+        ("jit_compile_false", False), ("jit_compile_true", True)
+    )
+    def test_model_compile_batched_ds(self, jit_compile):
+        def token_probability_fn(inputs):
+            prob = tf.constant([[0.0, 0.0, 0.0, 1.0]])
+            return tf.repeat(prob, 2, axis=0)
+
+        max_length = 5
+
+        class TestModel(tf.keras.Model):
+            def call(self, inputs):
+                generated = greedy_search(
+                    token_probability_fn,
+                    inputs,
+                    max_length=max_length,
+                    end_token_id=2,
+                    pad_token_id=0,
+                )
+                return generated
+
+        inputs = tf.constant([[0, 1], [1, 2]])
+        ds = tf.data.Dataset.from_tensor_slices(inputs).batch(2)
+        expected_outputs = tf.tile([[3], [0]], [1, max_length - 2])
+        expected_outputs = tf.concat([inputs, expected_outputs], axis=1)
+
+        model = TestModel()
+        model.compile(jit_compile=jit_compile)
+
+        outputs = model.predict(ds)
+        self.assertAllEqual(outputs, expected_outputs)
+
 
 class BeamSearchTextGenerationTest(tf.test.TestCase):
     def setUp(self):
@@ -249,7 +313,7 @@ def token_probability_fn(inputs):
     def test_assert_generation_is_correct(self):
         def token_probability_fn(inputs):
             batch_size = inputs.shape[0]
-            prob = tf.constant([[0.01, 0.01, 0.08, 0.9]])
+            prob = tf.constant([[0.0, 0.0, 0.0, 1.0]])
             return tf.repeat(prob, batch_size, axis=0)
 
         batch_size = 10
@@ -269,7 +333,7 @@ def token_probability_fn(inputs):
     def test_end_token_id(self):
         def token_probability_fn(inputs):
             batch_size = inputs.shape[0]
-            prob = tf.constant([[0.01, 0.01, 0.08, 0.9]])
+            prob = tf.constant([[0.0, 0.0, 0.0, 1.0]])
             return tf.repeat(prob, batch_size, axis=0)
 
         max_length = 5
@@ -287,7 +351,7 @@ def token_probability_fn(inputs):
         self.assertAllEqual(outputs, expected_outputs)
 
 
-class RandomSearchTextGenerationTest(tf.test.TestCase):
+class RandomSearchTextGenerationTest(tf.test.TestCase, parameterized.TestCase):
     def setUp(self):
         super().setUp()
         self.vocab_size = 10
@@ -342,7 +406,7 @@ def test_generate_with_ragged_prompt(self):
     def test_assert_probability_distribution_generation_is_correct(self):
         def token_probability_fn(inputs):
             batch_size = inputs.shape[0]
-            prob = tf.constant([[0.01, 0.01, 0.08, 0.9]])
+            prob = tf.constant([[0.0, 0.0, 0.0, 1.0]])
             return tf.repeat(prob, batch_size, axis=0)
 
         batch_size = 10
@@ -360,14 +424,14 @@ def token_probability_fn(inputs):
                 outputs_count[pred] += 1
         self.assertAllClose(
             outputs_count / np.sum(outputs_count),
-            [0.01, 0.01, 0.08, 0.9],
+            [0.0, 0.0, 0.0, 1.0],
             rtol=0.2,
         )
 
     def test_end_token_id(self):
         def token_probability_fn(inputs):
             batch_size = inputs.shape[0]
-            prob = tf.constant([[0.01, 0.01, 0.08, 0.9]])
+            prob = tf.constant([[0.0, 0.0, 0.0, 1.0]])
             return tf.repeat(prob, batch_size, axis=0)
 
         max_length = 5
@@ -386,6 +450,71 @@ def token_probability_fn(inputs):
         expected_outputs = tf.concat([inputs, expected_outputs], axis=1)
         self.assertAllEqual(outputs, expected_outputs)
 
+    @parameterized.named_parameters(
+        ("jit_compile_false", False), ("jit_compile_true", True)
+    )
+    def test_model_compile(self, jit_compile):
+        def token_probability_fn(inputs):
+            prob = tf.constant([[0.0, 0.0, 0.0, 1.0]])
+            return tf.repeat(prob, 2, axis=0)
+
+        max_length = 5
+
+        class TestModel(tf.keras.Model):
+            def call(self, inputs):
+                generated = random_search(
+                    token_probability_fn,
+                    inputs,
+                    max_length=max_length,
+                    seed=42,
+                    end_token_id=2,
+                    pad_token_id=0,
+                )
+                return generated
+
+        inputs = tf.constant([[0, 1], [1, 2]])
+        expected_outputs = tf.tile([[3], [0]], [1, max_length - 2])
+        expected_outputs = tf.concat([inputs, expected_outputs], axis=1)
+
+        model = TestModel()
+        model.compile(jit_compile=jit_compile)
+
+        outputs = model(inputs)
+        self.assertAllEqual(outputs, expected_outputs)
+
+    @parameterized.named_parameters(
+        ("jit_compile_false", False), ("jit_compile_true", True)
+    )
+    def test_model_compile_batched_ds(self, jit_compile):
+        def token_probability_fn(inputs):
+            prob = tf.constant([[0.0, 0.0, 0.0, 1.0]])
+            return tf.repeat(prob, 2, axis=0)
+
+        max_length = 5
+
+        class TestModel(tf.keras.Model):
+            def call(self, inputs):
+                generated = random_search(
+                    token_probability_fn,
+                    inputs,
+                    max_length=max_length,
+                    seed=42,
+                    end_token_id=2,
+                    pad_token_id=0,
+                )
+                return generated
+
+        inputs = tf.constant([[0, 1], [1, 2]])
+        ds = tf.data.Dataset.from_tensor_slices(inputs).batch(2)
+        expected_outputs = tf.tile([[3], [0]], [1, max_length - 2])
+        expected_outputs = tf.concat([inputs, expected_outputs], axis=1)
+
+        model = TestModel()
+        model.compile(jit_compile=jit_compile)
+
+        outputs = model.predict(ds)
+        self.assertAllEqual(outputs, expected_outputs)
+
     def test_from_logits(self):
         def token_logits_fn(inputs):
             batch_size = inputs.shape[0]
@@ -420,7 +549,7 @@ def token_probability_fn(inputs):
         self.assertAllEqual(output_logit, output_probs)
 
 
-class TopKSearchTextGenerationTest(tf.test.TestCase):
+class TopKSearchTextGenerationTest(tf.test.TestCase, parameterized.TestCase):
     def setUp(self):
         super().setUp()
         self.vocab_size = 10
@@ -491,7 +620,7 @@ def test_generate_with_ragged_prompt(self):
     def test_assert_probability_distribution_generation_is_correct(self):
         def token_probability_fn(inputs):
             batch_size = inputs.shape[0]
-            prob = tf.constant([[0.1, 0.2, 0.3, 0.4]])
+            prob = tf.constant([[0.0, 0.0, 0.0, 1.0]])
             return tf.repeat(prob, batch_size, axis=0)
 
         batch_size = 10
@@ -513,7 +642,7 @@ def token_probability_fn(inputs):
                 outputs_count[pred] += 1
         self.assertAllClose(
             outputs_count / np.sum(outputs_count),
-            [0.0, 0.0, 0.429, 0.571],
+            [0.0, 0.0, 0.0, 1.0],
             rtol=0.2,
         )
 
@@ -539,7 +668,7 @@ def token_probability_fn(inputs):
     def test_end_token_id(self):
         def token_probability_fn(inputs):
             batch_size = inputs.shape[0]
-            prob = tf.constant([[0.01, 0.01, 0.08, 0.9]])
+            prob = tf.constant([[0.0, 0.0, 0.0, 1.0]])
             return tf.repeat(prob, batch_size, axis=0)
 
         max_length = 5
@@ -559,6 +688,71 @@ def token_probability_fn(inputs):
         expected_outputs = tf.concat([inputs, expected_outputs], axis=1)
         self.assertAllEqual(outputs, expected_outputs)
 
+    @parameterized.named_parameters(
+        ("jit_compile_false", False), ("jit_compile_true", True)
+    )
+    def test_model_compile(self, jit_compile):
+        def token_probability_fn(inputs):
+            prob = tf.constant([[0.0, 0.0, 0.0, 1.0]])
+            return tf.repeat(prob, 2, axis=0)
+
+        max_length = 5
+
+        class TestModel(tf.keras.Model):
+            def call(self, inputs):
+                generated = random_search(
+                    token_probability_fn,
+                    inputs,
+                    max_length=max_length,
+                    seed=42,
+                    end_token_id=2,
+                    pad_token_id=0,
+                )
+                return generated
+
+        inputs = tf.constant([[0, 1], [1, 2]])
+        expected_outputs = tf.tile([[3], [0]], [1, max_length - 2])
+        expected_outputs = tf.concat([inputs, expected_outputs], axis=1)
+
+        model = TestModel()
+        model.compile(jit_compile=jit_compile)
+
+        outputs = model.predict(inputs)
+        self.assertAllEqual(outputs, expected_outputs)
+
+    @parameterized.named_parameters(
+        ("jit_compile_false", False), ("jit_compile_true", True)
+    )
+    def test_model_compile_batched_ds(self, jit_compile):
+        def token_probability_fn(inputs):
+            prob = tf.constant([[0.0, 0.0, 0.0, 1.0]])
+            return tf.repeat(prob, 2, axis=0)
+
+        max_length = 5
+
+        class TestModel(tf.keras.Model):
+            def call(self, inputs):
+                generated = random_search(
+                    token_probability_fn,
+                    inputs,
+                    max_length=max_length,
+                    seed=42,
+                    end_token_id=2,
+                    pad_token_id=0,
+                )
+                return generated
+
+        inputs = tf.constant([[0, 1], [1, 2]])
+        ds = tf.data.Dataset.from_tensor_slices(inputs).batch(2)
+        expected_outputs = tf.tile([[3], [0]], [1, max_length - 2])
+        expected_outputs = tf.concat([inputs, expected_outputs], axis=1)
+
+        model = TestModel()
+        model.compile(jit_compile=jit_compile)
+
+        outputs = model.predict(ds)
+        self.assertAllEqual(outputs, expected_outputs)
+
     def test_from_logits(self):
         def token_logits_fn(inputs):
             batch_size = inputs.shape[0]
@@ -595,7 +789,7 @@ def token_probability_fn(inputs):
         self.assertAllEqual(output_logit, output_probs)
 
 
-class TopPSearchTextGenerationTest(tf.test.TestCase):
+class TopPSearchTextGenerationTest(tf.test.TestCase, parameterized.TestCase):
     def setUp(self):
         super().setUp()
         self.vocab_size = 10
@@ -656,7 +850,7 @@ def test_generate_with_ragged_prompt(self):
     def test_assert_probability_distribution_generation_is_correct(self):
         def token_probability_fn(inputs):
             batch_size = inputs.shape[0]
-            prob = tf.constant([[0.1, 0.2, 0.3, 0.4]])
+            prob = tf.constant([[0.0, 0.0, 0.0, 1.0]])
             return tf.repeat(prob, batch_size, axis=0)
 
         batch_size = 10
@@ -678,7 +872,7 @@ def token_probability_fn(inputs):
                 outputs_count[pred] += 1
         self.assertAllClose(
             outputs_count / np.sum(outputs_count),
-            [0.0, 0.0, 0.429, 0.571],
+            [0.0, 0.0, 0.0, 1.0],
             rtol=0.2,
         )
 
@@ -700,7 +894,7 @@ def token_probability_fn(inputs):
     def test_end_token_id(self):
         def token_probability_fn(inputs):
             batch_size = inputs.shape[0]
-            prob = tf.constant([[0.01, 0.01, 0.08, 0.9]])
+            prob = tf.constant([[0.0, 0.0, 0.0, 1.0]])
             return tf.repeat(prob, batch_size, axis=0)
 
         max_length = 5
@@ -720,6 +914,73 @@ def token_probability_fn(inputs):
         expected_outputs = tf.concat([inputs, expected_outputs], axis=1)
         self.assertAllEqual(outputs, expected_outputs)
 
+    @parameterized.named_parameters(
+        ("jit_compile_false", False), ("jit_compile_true", True)
+    )
+    def test_model_compile(self, jit_compile):
+        def token_probability_fn(inputs):
+            prob = tf.constant([[0.0, 0.0, 0.0, 1.0]])
+            return tf.repeat(prob, 2, axis=0)
+
+        max_length = 5
+
+        class TestModel(tf.keras.Model):
+            def call(self, inputs):
+                generated = top_p_search(
+                    token_probability_fn,
+                    inputs,
+                    max_length=max_length,
+                    p=0.92,
+                    seed=1,
+                    end_token_id=2,
+                    pad_token_id=0,
+                )
+                return generated
+
+        inputs = tf.constant([[0, 1], [1, 2]])
+        expected_outputs = tf.tile([[3], [0]], [1, max_length - 2])
+        expected_outputs = tf.concat([inputs, expected_outputs], axis=1)
+
+        model = TestModel()
+        model.compile(jit_compile=jit_compile)
+
+        outputs = model.predict(inputs)
+        self.assertAllEqual(outputs, expected_outputs)
+
+    @parameterized.named_parameters(
+        ("jit_compile_false", False), ("jit_compile_true", True)
+    )
+    def test_model_compile_batched_ds(self, jit_compile):
+        def token_probability_fn(inputs):
+            prob = tf.constant([[0.0, 0.0, 0.0, 1.0]])
+            return tf.repeat(prob, 2, axis=0)
+
+        max_length = 5
+
+        class TestModel(tf.keras.Model):
+            def call(self, inputs):
+                generated = top_p_search(
+                    token_probability_fn,
+                    inputs,
+                    max_length=max_length,
+                    p=0.92,
+                    seed=1,
+                    end_token_id=2,
+                    pad_token_id=0,
+                )
+                return generated
+
+        inputs = tf.constant([[0, 1], [1, 2]])
+        ds = tf.data.Dataset.from_tensor_slices(inputs).batch(2)
+        expected_outputs = tf.tile([[3], [0]], [1, max_length - 2])
+        expected_outputs = tf.concat([inputs, expected_outputs], axis=1)
+
+        model = TestModel()
+        model.compile(jit_compile=jit_compile)
+
+        outputs = model.predict(ds)
+        self.assertAllEqual(outputs, expected_outputs)
+
     def test_from_logits(self):
         def token_logits_fn(inputs):
             batch_size = inputs.shape[0]