Repository: keras-team/keras-nlp
Pull request #1186 (2023-07-28 16:59:50)

Port metrics to multi-backend
Probably more we could do here, but here is a quick port of our metrics to multi-backend.

In short, we will leave them as eager-only metrics, and all "string metrics" will use `tf.Tensor` as an intermediate data type. This is really cross-platform offering that makes sense for these metrics as we lack string types in other platforms. It's not a huge loss either, as most of these metrics are things you want to apply to generated sequences from a model, so you could never really bake most of these in as "live" metrics computed during a `train_step`.

The one exception here is perplexity. This metric is fully in integer space, and we can write it entirely with out `ops` layer, and compile it across platforms.

mattdangerw: /gcbrun

Patch:
diff --git a/keras_nlp/layers/modeling/masked_lm_head.py b/keras_nlp/layers/modeling/masked_lm_head.py
index 0bb0a421..71014c47 100644
--- a/keras_nlp/layers/modeling/masked_lm_head.py
+++ b/keras_nlp/layers/modeling/masked_lm_head.py
@@ -123,7 +123,6 @@ class MaskedLMHead(keras.layers.Layer):
         self.layer_norm_epsilon = layer_norm_epsilon
         self.kernel_initializer = keras.initializers.get(kernel_initializer)
         self.bias_initializer = keras.initializers.get(bias_initializer)
-        self._built = False
 
         if vocabulary_size is None and embedding_weights is None:
             raise ValueError(
@@ -142,7 +141,7 @@ class MaskedLMHead(keras.layers.Layer):
                 )
             self.vocabulary_size = shape[0]
 
-    def build(self, inputs_shape, masked_positions_shape=None):
+    def build(self, inputs_shape):
         if self.embedding_weights is not None:
             feature_size = self.embedding_weights.shape[-1]
         else:
@@ -157,12 +156,13 @@ class MaskedLMHead(keras.layers.Layer):
         self._layer_norm = keras.layers.LayerNormalization(
             epsilon=self.layer_norm_epsilon,
         )
-        if masked_positions_shape:
-            gather_length = masked_positions_shape[1]
-            shape = (inputs_shape[0], gather_length, inputs_shape[-1])
-            self._dense.build(shape)
-            shape = (inputs_shape[0], gather_length, feature_size)
-            self._layer_norm.build(shape)
+        # The gather length does not affect any of our built variables, so
+        # we can pass any value here.
+        gather_length = None
+        shape = (inputs_shape[0], gather_length, inputs_shape[-1])
+        self._dense.build(shape)
+        shape = (inputs_shape[0], gather_length, feature_size)
+        self._layer_norm.build(shape)
         if self.embedding_weights is None:
             self._kernel = self.add_weight(
                 name="output_kernel",
@@ -177,10 +177,10 @@ class MaskedLMHead(keras.layers.Layer):
             dtype=self.dtype,
         )
 
-    def call(self, inputs, masked_positions):
+    def call(self, inputs, mask_positions):
         # Gather the encoded tokens at the masked indices.
-        masked_positions = ops.expand_dims(masked_positions, axis=-1)
-        x = ops.take_along_axis(inputs, masked_positions, axis=1)
+        mask_positions = ops.expand_dims(mask_positions, axis=-1)
+        x = ops.take_along_axis(inputs, mask_positions, axis=1)
 
         # Apply a trainable linear transformation and a layer norm.
         x = self._dense(x)
@@ -221,7 +221,9 @@ class MaskedLMHead(keras.layers.Layer):
         )
         return config
 
-    def compute_output_shape(self, inputs_shape, masked_positions_shape):
-        output_shape = list(masked_positions_shape)
-        output_shape[-1] = self.vocabulary_size
-        return tuple(output_shape)
+    # TODO: restore this after https://github.com/keras-team/keras-core/pull/632
+    # is in a release!
+    # def compute_output_shape(self, inputs_shape, mask_positions_shape):
+    #     output_shape = list(mask_positions_shape)
+    #     output_shape[-1] = self.vocabulary_size
+    #     return tuple(output_shape)
diff --git a/keras_nlp/metrics/bleu.py b/keras_nlp/metrics/bleu.py
index e1029f0d..3ee35dda 100644
--- a/keras_nlp/metrics/bleu.py
+++ b/keras_nlp/metrics/bleu.py
@@ -21,7 +21,7 @@ import tensorflow as tf
 
 from keras_nlp.api_export import keras_nlp_export
 from keras_nlp.backend import keras
-from keras_nlp.utils.tensor_utils import assert_tf_backend
+from keras_nlp.backend import ops
 from keras_nlp.utils.tensor_utils import is_floating_dtype
 from keras_nlp.utils.tensor_utils import tensor_to_list
 
@@ -112,8 +112,6 @@ class Bleu(keras.metrics.Metric):
         name="bleu",
         **kwargs,
     ):
-        assert_tf_backend(self.__class__.__name__)
-
         super().__init__(name=name, dtype=dtype, **kwargs)
 
         if not is_floating_dtype(dtype):
@@ -290,8 +288,10 @@ class Bleu(keras.metrics.Metric):
         )
 
     def _calculate_bleu_score(self, references, translation):
-        references = tensor_to_list(references)
-        translation = tensor_to_list(translation)
+        if isinstance(references, (tf.Tensor, tf.RaggedTensor)):
+            references = tensor_to_list(references)
+        if isinstance(translation, (tf.Tensor, tf.RaggedTensor)):
+            translation = tensor_to_list(translation)
 
         matches = self._matches.numpy()
         possible_matches = self._possible_matches.numpy()
@@ -315,11 +315,11 @@ class Bleu(keras.metrics.Metric):
             smooth=self.smooth,
         )
         return (
-            tf.constant(bleu_score, dtype=self.dtype),
-            tf.constant(matches, dtype=self.dtype),
-            tf.constant(possible_matches, dtype=self.dtype),
-            tf.constant(translation_length, dtype=self.dtype),
-            tf.constant(reference_length, dtype=self.dtype),
+            bleu_score,
+            matches,
+            possible_matches,
+            translation_length,
+            reference_length,
         )
 
     def update_state(self, y_true, y_pred, sample_weight=None):
@@ -357,11 +357,7 @@ class Bleu(keras.metrics.Metric):
             possible_matches,
             translation_length,
             reference_length,
-        ) = tf.py_function(
-            func=self._calculate_bleu_score,
-            inp=[y_true, y_pred],
-            Tout=[self.dtype, self.dtype, self.dtype, self.dtype, self.dtype],
-        )
+        ) = self._calculate_bleu_score(y_true, y_pred)
 
         self._matches.assign(matches)
         self._possible_matches.assign(possible_matches)
@@ -374,10 +370,10 @@ class Bleu(keras.metrics.Metric):
 
     def reset_state(self):
         self._matches.assign(
-            tf.zeros(shape=(self.max_order,), dtype=self.dtype)
+            ops.zeros(shape=(self.max_order,), dtype=self.dtype)
         )
         self._possible_matches.assign(
-            tf.zeros(shape=(self.max_order,), dtype=self.dtype)
+            ops.zeros(shape=(self.max_order,), dtype=self.dtype)
         )
         self._translation_length.assign(0.0)
         self._reference_length.assign(0.0)
diff --git a/keras_nlp/metrics/edit_distance.py b/keras_nlp/metrics/edit_distance.py
index c1ab7ba2..fd41bde8 100644
--- a/keras_nlp/metrics/edit_distance.py
+++ b/keras_nlp/metrics/edit_distance.py
@@ -18,7 +18,6 @@ import tensorflow as tf
 
 from keras_nlp.api_export import keras_nlp_export
 from keras_nlp.backend import keras
-from keras_nlp.utils.tensor_utils import assert_tf_backend
 from keras_nlp.utils.tensor_utils import is_floating_dtype
 
 
@@ -67,13 +66,6 @@ class EditDistance(keras.metrics.Metric):
     >>> edit_distance(y_true, y_pred)
     <tf.Tensor: shape=(), dtype=float32, numpy=0.36363637>
 
-    Rank 1 tensor.
-    >>> edit_distance = keras_nlp.metrics.EditDistance()
-    >>> y_true = tf.strings.split("the tiny little cat was found under the big funny bed")
-    >>> y_pred = tf.strings.split("the cat was found under the bed")
-    >>> edit_distance(y_true, y_pred)
-    <tf.Tensor: shape=(), dtype=float32, numpy=0.36363637>
-
     Nested Python list.
     >>> edit_distance = keras_nlp.metrics.EditDistance()
     >>> y_true = [
@@ -86,19 +78,6 @@ class EditDistance(keras.metrics.Metric):
     ... ]
     >>> edit_distance(y_true, y_pred)
     <tf.Tensor: shape=(), dtype=float32, numpy=0.73333335>
-
-    Rank 2 tensor.
-    >>> edit_distance = keras_nlp.metrics.EditDistance()
-    >>> y_true = tf.strings.split([
-    ...     "the tiny little cat was found under the big funny bed",
-    ...     "it is sunny today",
-    ... ])
-    >>> y_pred = tf.strings.split([
-    ...     "the cat was found under the bed",
-    ...     "it is sunny but with a hint of cloud cover",
-    ... ])
-    >>> edit_distance(y_true, y_pred)
-    <tf.Tensor: shape=(), dtype=float32, numpy=0.73333335>
     """
 
     def __init__(
@@ -108,8 +87,6 @@ class EditDistance(keras.metrics.Metric):
         name="edit_distance",
         **kwargs,
     ):
-        assert_tf_backend(self.__class__.__name__)
-
         super().__init__(name=name, dtype=dtype, **kwargs)
 
         if not is_floating_dtype(dtype):
diff --git a/keras_nlp/metrics/perplexity.py b/keras_nlp/metrics/perplexity.py
index 1a2398fc..7ff78ebc 100644
--- a/keras_nlp/metrics/perplexity.py
+++ b/keras_nlp/metrics/perplexity.py
@@ -14,11 +14,9 @@
 
 """Perplexity metric."""
 
-import tensorflow as tf
-
 from keras_nlp.api_export import keras_nlp_export
 from keras_nlp.backend import keras
-from keras_nlp.utils.tensor_utils import assert_tf_backend
+from keras_nlp.backend import ops
 from keras_nlp.utils.tensor_utils import is_floating_dtype
 
 
@@ -48,46 +46,40 @@ class Perplexity(keras.metrics.Metric):
 
     1. Calculate perplexity by calling update_state() and result().
     1.1. `sample_weight`, and `mask_token_id` are not provided.
-    >>> tf.random.set_seed(42)
+    >>> np.random.seed(42)
     >>> perplexity = keras_nlp.metrics.Perplexity(name="perplexity")
-    >>> target = tf.random.uniform(
-    ...     shape=[2, 5],  maxval=10, dtype="int32", seed=42)
-    >>> logits = tf.random.uniform(shape=(2, 5, 10), seed=42)
+    >>> target = np.random.randint(10, size=[2, 5])
+    >>> logits = np.random.uniform(size=(2, 5, 10))
     >>> perplexity.update_state(target, logits)
     >>> perplexity.result()
-    <tf.Tensor: shape=(), dtype=float32, numpy=11.8781595>
+    <tf.Tensor: shape=(), dtype=float32, numpy=14.352535>
 
     1.2. `sample_weight` specified (masking token with ID 0).
-    >>> tf.random.set_seed(42)
+    >>> np.random.seed(42)
     >>> perplexity = keras_nlp.metrics.Perplexity(name="perplexity")
-    >>> target = tf.random.uniform(
-    ...     shape=[2, 5],  maxval=10, dtype="int32", seed=42)
-    >>> logits = tf.random.uniform(shape=(2, 5, 10), seed=42)
-    >>> sample_weight = tf.cast(
-    ...     tf.math.logical_not(tf.equal(target, 0)), "float32")
+    >>> target = np.random.randint(10, size=[2, 5])
+    >>> logits = np.random.uniform(size=(2, 5, 10))
+    >>> sample_weight = (target != 0).astype("float32")
     >>> perplexity.update_state(target, logits, sample_weight)
     >>> perplexity.result()
-    <tf.Tensor: shape=(), dtype=float32, numpy=13.1128>
+    <tf.Tensor: shape=(), dtype=float32, numpy=14.352535>
 
     2. Call perplexity directly.
-    >>> tf.random.set_seed(42)
+    >>> np.random.seed(42)
     >>> perplexity = keras_nlp.metrics.Perplexity(name="perplexity")
-    >>> target = tf.random.uniform(
-    ...     shape=[2, 5],  maxval=10, dtype="int32", seed=42)
-    >>> logits = tf.random.uniform(shape=(2, 5, 10), seed=42)
+    >>> target = np.random.randint(10, size=[2, 5])
+    >>> logits = np.random.uniform(size=(2, 5, 10))
     >>> perplexity(target, logits)
-    <tf.Tensor: shape=(), dtype=float32, numpy=11.8781595>
+    <tf.Tensor: shape=(), dtype=float32, numpy=14.352535>
 
     3. Provide the padding token ID and let the class compute the mask on its
        own.
-    >>> tf.random.set_seed(42)
-    >>> perplexity = keras_nlp.metrics.Perplexity(
-    ...     name="perplexity", mask_token_id=0)
-    >>> target = tf.random.uniform(
-    ...     shape=[2, 5],  maxval=10, dtype="int32", seed=42)
-    >>> logits = tf.random.uniform(shape=(2, 5, 10), seed=42)
+    >>> np.random.seed(42)
+    >>> perplexity = keras_nlp.metrics.Perplexity(mask_token_id=0)
+    >>> target = np.random.randint(10, size=[2, 5])
+    >>> logits = np.random.uniform(size=(2, 5, 10))
     >>> perplexity(target, logits)
-    <tf.Tensor: shape=(), dtype=float32, numpy=13.1128>
+    <tf.Tensor: shape=(), dtype=float32, numpy=14.352535>
     """
 
     def __init__(
@@ -98,8 +90,6 @@ class Perplexity(keras.metrics.Metric):
         name="perplexity",
         **kwargs,
     ):
-        assert_tf_backend(self.__class__.__name__)
-
         if not is_floating_dtype(dtype):
             raise ValueError(
                 "`dtype` must be a floating point type. "
@@ -131,39 +121,39 @@ class Perplexity(keras.metrics.Metric):
     def update_state(self, y_true, y_pred, sample_weight=None):
         # y_true shape: (batch_size, seq_len)
         # y_pred shape: (batch_size, seq_len, vocab_size)
-        y_true = tf.cast(y_true, self.dtype)
-        y_pred = tf.cast(y_pred, self.dtype)
+        y_true = ops.cast(y_true, self.dtype)
+        y_pred = ops.cast(y_pred, self.dtype)
 
         if sample_weight is not None:
-            sample_weight = tf.cast(sample_weight, self.dtype)
+            sample_weight = ops.cast(sample_weight, self.dtype)
 
-        batch_size = tf.cast(tf.shape(y_true)[0], self.dtype)
+        batch_size = ops.cast(ops.shape(y_true)[0], self.dtype)
 
         if self.mask_token_id is not None:
-            mask = tf.cast(
-                tf.math.logical_not(tf.equal(y_true, self.mask_token_id)),
+            mask = ops.cast(
+                ops.logical_not(ops.equal(y_true, self.mask_token_id)),
                 self.dtype,
             )
             if sample_weight is None:
                 sample_weight = mask
             else:
-                sample_weight = tf.multiply(mask, sample_weight)
+                sample_weight = ops.multiply(mask, sample_weight)
 
         # Calculate the Cross Entropy Loss.
-        crossentropy_value = tf.cast(
+        crossentropy_value = ops.cast(
             self._crossentropy(y_true, y_pred, sample_weight=sample_weight),
             self.dtype,
         )  # scalar
 
         # Divide the loss by the number of non-masked tokens
         if sample_weight is not None:
-            crossentropy_value = crossentropy_value / tf.reduce_sum(
+            crossentropy_value = crossentropy_value / ops.sum(
                 sample_weight
             )  # scalar
         else:
             crossentropy_value = crossentropy_value / (
-                tf.cast(tf.shape(y_true)[0], self.dtype)
-                * tf.cast(tf.shape(y_true)[1], self.dtype)
+                ops.cast(ops.shape(y_true)[0], self.dtype)
+                * ops.cast(ops.shape(y_true)[1], self.dtype)
             )  # scalar
 
         self._aggregate_crossentropy.assign_add(batch_size * crossentropy_value)
@@ -172,7 +162,7 @@ class Perplexity(keras.metrics.Metric):
     def result(self):
         if self._number_of_samples == 0:
             return 0.0
-        perplexity_score = tf.exp(
+        perplexity_score = ops.exp(
             self._aggregate_crossentropy / self._number_of_samples
         )
         return perplexity_score
diff --git a/keras_nlp/metrics/rouge_base.py b/keras_nlp/metrics/rouge_base.py
index aba2df05..deec9536 100644
--- a/keras_nlp/metrics/rouge_base.py
+++ b/keras_nlp/metrics/rouge_base.py
@@ -15,12 +15,10 @@
 """ROUGE metric."""
 
 
-import types
-
 import tensorflow as tf
 
 from keras_nlp.backend import keras
-from keras_nlp.utils.tensor_utils import assert_tf_backend
+from keras_nlp.backend import ops
 from keras_nlp.utils.tensor_utils import is_floating_dtype
 from keras_nlp.utils.tensor_utils import tensor_to_list
 
@@ -62,8 +60,6 @@ class RougeBase(keras.metrics.Metric):
         name="rouge",
         **kwargs,
     ):
-        assert_tf_backend(self.__class__.__name__)
-
         super().__init__(name=name, dtype=dtype, **kwargs)
 
         if rouge_scorer is None:
@@ -123,27 +119,6 @@ class RougeBase(keras.metrics.Metric):
             name="number_of_samples",
         )
 
-    def __new__(cls, *args, **kwargs):
-        # Temporary workaround for Keras bug with dictionary return types.
-        # Wraps `result()` with a python dictionary that also supports variable
-        # assignment. We have to do this with __new__ because the base metric
-        # class wraps the `results()` method.
-        # TODO: Remove this snippet of code once the Keras bug is fixed.
-        obj = super().__new__(cls)
-
-        class MetricDict(dict):
-            """A dictionary that supports variable assignment."""
-
-            pass
-
-        def wrap_result(result_fn):
-            return tf.__internal__.decorator.make_decorator(
-                result_fn, lambda obj, *args: MetricDict(result_fn(*args))
-            )
-
-        obj.result = types.MethodType(wrap_result(obj.result), obj)
-        return obj
-
     def update_state(self, y_true, y_pred, sample_weight=None):
         # Three possible shapes for y_true and y_pred: Python string,
         # [batch_size] and [batch_size, 1]. In the latter two cases, we have
@@ -182,23 +157,16 @@ class RougeBase(keras.metrics.Metric):
             score = self._rouge_scorer.score(reference, hypothesis)[
                 self.variant
             ]
-            return tf.cast(
-                tf.constant([score.precision, score.recall, score.fmeasure]),
-                dtype=self.dtype,
-            )
+            return score.precision, score.recall, score.fmeasure
 
         for batch_idx in range(batch_size):
-            score = tf.py_function(
-                func=calculate_rouge_score,
-                inp=[y_true[batch_idx], y_pred[batch_idx]],
-                Tout=self.dtype,
-            )
+            score = calculate_rouge_score(y_true[batch_idx], y_pred[batch_idx])
             self._rouge_precision.assign_add(score[0])
             self._rouge_recall.assign_add(score[1])
             self._rouge_f1_score.assign_add(score[2])
 
         self._number_of_samples.assign_add(
-            tf.cast(batch_size, dtype=self.dtype)
+            ops.cast(batch_size, dtype=self.dtype)
         )
 
     def result(self):
diff --git a/keras_nlp/metrics/rouge_l.py b/keras_nlp/metrics/rouge_l.py
index ffa62bbb..216cd6ed 100644
--- a/keras_nlp/metrics/rouge_l.py
+++ b/keras_nlp/metrics/rouge_l.py
@@ -44,15 +44,14 @@ class RougeL(RougeBase):
 
     Examples:
 
-    1. Various Input Types.
-    1.1. Python string.
+    1. Python string.
     >>> rouge_l = keras_nlp.metrics.RougeL()
     >>> y_true = "the tiny little cat was found under the big funny bed"
     >>> y_pred = "the cat was under the bed"
     >>> rouge_l(y_true, y_pred)["f1_score"]
     <tf.Tensor: shape=(), dtype=float32, numpy=0.7058824>
 
-    1.2. rank 1 inputs.
+    2. List inputs.
     a. Python list.
     >>> rouge_l = keras_nlp.metrics.RougeL()
     >>> y_true = [
@@ -66,49 +65,19 @@ class RougeL(RougeBase):
     >>> rouge_l(y_true, y_pred)["f1_score"]
     <tf.Tensor: shape=(), dtype=float32, numpy=0.80748665>
 
-    b. Tensor
-    >>> rouge_l = keras_nlp.metrics.RougeL()
-    >>> y_true = tf.constant(
-    ...     [
-    ...         "the tiny little cat was found under the big funny bed",
-    ...         "i really love contributing to KerasNLP",
-    ...     ]
-    ... )
-    >>> y_pred = tf.constant(
-    ...     [
-    ...         "the cat was under the bed",
-    ...         "i love contributing to KerasNLP",
-    ...     ]
-    ... )
-    >>> rouge_l(y_true, y_pred)["f1_score"]
-    <tf.Tensor: shape=(), dtype=float32, numpy=0.80748665>
 
-    1.3. rank 2 inputs.
+    3. 2D inputs.
     >>> rouge_l = keras_nlp.metrics.RougeL()
-    >>> y_true = tf.constant(
-    ...     [
-    ...         ["the tiny little cat was found under the big funny bed"],
-    ...         ["i really love contributing to KerasNLP"],
-    ...     ]
-    ... )
-    >>> y_pred = tf.constant(
-    ...     [
-    ...         ["the cat was under the bed"],
-    ...         ["i love contributing to KerasNLP"],
-    ...     ]
-    ... )
+    >>> y_true = [
+    ...     ["the tiny little cat was found under the big funny bed"],
+    ...     ["i really love contributing to KerasNLP"],
+    ... ]
+    >>> y_pred = [
+    ...     ["the cat was under the bed"],
+    ...     ["i love contributing to KerasNLP"],
+    ... ]
     >>> rouge_l(y_true, y_pred)["f1_score"]
     <tf.Tensor: shape=(), dtype=float32, numpy=0.80748665>
-
-    3. Pass the metric to `model.compile()`.
-    >>> inputs = keras.Input(shape=(), dtype='string')
-    >>> outputs = keras.layers.Identity()(inputs)
-    >>> model = keras.Model(inputs, outputs)
-    >>> model.compile(metrics=[keras_nlp.metrics.RougeL()])
-    >>> y_pred = x = tf.constant(["hello this is fun"])
-    >>> y = tf.constant(["hello this is awesome"])
-    >>> model.compute_metrics(x, y, y_pred, sample_weight=None)["f1_score"]
-    0.75
     """
 
     def __init__(
diff --git a/keras_nlp/metrics/rouge_n.py b/keras_nlp/metrics/rouge_n.py
index d90122fa..77bfd702 100644
--- a/keras_nlp/metrics/rouge_n.py
+++ b/keras_nlp/metrics/rouge_n.py
@@ -46,16 +46,14 @@ class RougeN(RougeBase):
 
     Examples:
 
-    1. Various Input Types.
-    1.1. Python string.
+    1. Python string.
     >>> rouge_n = keras_nlp.metrics.RougeN(order=2)
     >>> y_true = "the tiny little cat was found under the big funny bed"
     >>> y_pred = "the cat was under the bed"
     >>> rouge_n(y_true, y_pred)["f1_score"]
     <tf.Tensor: shape=(), dtype=float32, numpy=0.26666668>
 
-    1.2. rank 1 inputs.
-    a. Python list.
+    2. List inputs.
     >>> rouge_n = keras_nlp.metrics.RougeN(order=2)
     >>> y_true = [
     ...     "the tiny little cat was found under the big funny bed",
@@ -68,66 +66,31 @@ class RougeN(RougeBase):
     >>> rouge_n(y_true, y_pred)["f1_score"]
     <tf.Tensor: shape=(), dtype=float32, numpy=0.4666667>
 
-    b. Tensor.
+    3. 2D inputs.
     >>> rouge_n = keras_nlp.metrics.RougeN(order=2)
-    >>> y_true = tf.constant(
-    ...     [
-    ...         "the tiny little cat was found under the big funny bed",
-    ...         "i really love contributing to KerasNLP",
-    ...     ]
-    ... )
-    >>> y_pred = tf.constant(
-    ...     [
-    ...         "the cat was under the bed",
-    ...         "i love contributing to KerasNLP",
-    ...     ]
-    ... )
-    >>> rouge_n(y_true, y_pred)["f1_score"]
-    <tf.Tensor: shape=(), dtype=float32, numpy=0.4666667>
-
-    1.3. rank 2 inputs.
-    >>> rouge_n = keras_nlp.metrics.RougeN(order=2)
-    >>> y_true = tf.constant(
-    ...     [
-    ...         ["the tiny little cat was found under the big funny bed"],
-    ...         ["i really love contributing to KerasNLP"],
-    ...     ]
-    ... )
-    >>> y_pred = tf.constant(
-    ...     [
-    ...         ["the cat was under the bed"],
-    ...         ["i love contributing to KerasNLP"],
-    ...     ]
-    ... )
+    >>> y_true =[
+    ...     ["the tiny little cat was found under the big funny bed"],
+    ...     ["i really love contributing to KerasNLP"],
+    ... ]
+    >>> y_pred =[
+    ...     ["the cat was under the bed"],
+    ...     ["i love contributing to KerasNLP"],
+    ... ]
     >>> rouge_n(y_true, y_pred)["f1_score"]
     <tf.Tensor: shape=(), dtype=float32, numpy=0.4666667>
 
-    2. Consider trigrams for calculating ROUGE-N.
+    4. Trigrams.
     >>> rouge_n = keras_nlp.metrics.RougeN(order=3)
-    >>> y_true = tf.constant(
-    ...     [
-    ...         "the tiny little cat was found under the big funny bed",
-    ...         "i really love contributing to KerasNLP",
-    ...     ]
-    ... )
-    >>> y_pred = tf.constant(
-    ...     [
-    ...         "the cat was under the bed",
-    ...         "i love contributing to KerasNLP",
-    ...     ]
-    ... )
+    >>> y_true = [
+    ...     "the tiny little cat was found under the big funny bed",
+    ...     "i really love contributing to KerasNLP",
+    ... ]
+    >>> y_pred = [
+    ...     "the cat was under the bed",
+    ...     "i love contributing to KerasNLP",
+    ... ]
     >>> rouge_n(y_true, y_pred)["f1_score"]
     <tf.Tensor: shape=(), dtype=float32, numpy=0.2857143>
-
-    3. Pass the metric to `model.compile()`.
-    >>> inputs = keras.Input(shape=(), dtype='string')
-    >>> outputs = keras.layers.Identity()(inputs)
-    >>> model = keras.Model(inputs, outputs)
-    >>> model.compile(metrics=[keras_nlp.metrics.RougeN()])
-    >>> y_pred = x = tf.constant(["hello this is fun"])
-    >>> y = tf.constant(["hello this is awesome"])
-    >>> model.compute_metrics(x, y, y_pred, sample_weight=None)["f1_score"]
-    0.6666666865348816
     """
 
     def __init__(
diff --git a/keras_nlp/samplers/beam_sampler.py b/keras_nlp/samplers/beam_sampler.py
index 077f250b..6a076f35 100644
--- a/keras_nlp/samplers/beam_sampler.py
+++ b/keras_nlp/samplers/beam_sampler.py
@@ -14,14 +14,13 @@
 """Beam Sampler."""
 
 import tensorflow as tf
-from tensorflow.compiler.tf2xla.python.xla import dynamic_update_slice
 
 from keras_nlp.api_export import keras_nlp_export
 from keras_nlp.backend import keras
+from keras_nlp.backend import ops
 from keras_nlp.samplers.sampler import Sampler
 from keras_nlp.samplers.sampler import call_args_docstring
 from keras_nlp.utils.python_utils import format_docstring
-from keras_nlp.utils.tensor_utils import assert_tf_backend
 
 
 @format_docstring(call_args=call_args_docstring)
@@ -102,10 +101,6 @@ class BeamSampler(Sampler):
         return_all_beams=False,
         **kwargs,
     ):
-        # Temporarily turn off beam search in other backends.
-        # No technical blockers here, just need tf -> ops rewrite.
-        assert_tf_backend(self.__class__.__name__)
-
         super().__init__(**kwargs)
         self.num_beams = num_beams
         self.return_all_beams = return_all_beams
@@ -120,89 +115,96 @@ class BeamSampler(Sampler):
         end_token_id=None,
         hidden_states=None,
     ):
-        batch_size, max_length = tf.shape(prompt)[0], tf.shape(prompt)[1]
-        # Make sure max length and start index are the same dtype.
-        index = tf.cast(index, max_length.dtype)
+        batch_size, max_length = ops.shape(prompt)[0], ops.shape(prompt)[1]
+        index = ops.cast(index, "int32")
 
         def create_beams(x):
             """Add initial beam state."""
-            return tf.repeat(x, self.num_beams, axis=0)
+            return ops.repeat(x, self.num_beams, axis=0)
 
         def flatten_beams(x):
             """Combine the beam dim and batch dim."""
-            flat_shape = [batch_size * self.num_beams] + x.shape.as_list()[2:]
-            return tf.reshape(x, shape=flat_shape)
+            flat_shape = (batch_size * self.num_beams,) + tuple(x.shape)[2:]
+            return ops.reshape(x, flat_shape)
 
         def unflatten_beams(x):
             """Separate the beam dim and batch dim."""
-            unflat_shape = [batch_size, self.num_beams] + x.shape.as_list()[1:]
-            return tf.reshape(x, shape=unflat_shape)
+            unflat_shape = (batch_size, self.num_beams) + tuple(x.shape)[1:]
+            return ops.reshape(x, unflat_shape)
 
         if mask is None:
-            mask = tf.zeros_like(prompt, dtype="bool")
+            mask = ops.zeros_like(prompt, dtype="bool")
         else:
-            mask = tf.cast(mask, dtype="bool")
-        # `tf.while_loop` will not accept `None` as a value for `loop_vars`.
-        cache = () if cache is None else cache
+            mask = ops.cast(mask, dtype="bool")
+        # `ops.while_loop` will not accept `None` as a value for `loop_vars`.
+        has_cache = cache is not None
+        cache = cache if has_cache else ()
         # Add extra sequences for each beam.
         prompt, mask = create_beams(prompt), create_beams(mask)
         cache = tf.nest.map_structure(create_beams, cache)
         # Setup the initial beam log-likelihoods.
         # On the first loop, make sure only the original beam is considered.
-        log_probs = tf.constant([[0.0] + [-1e9] * (self.num_beams - 1)])
-        log_probs = flatten_beams(tf.repeat(log_probs, batch_size, axis=0))
+        log_probs = ops.array(
+            [[0.0] + [-1e9] * (self.num_beams - 1)], dtype="float32"
+        )
+        log_probs = flatten_beams(ops.repeat(log_probs, batch_size, axis=0))
 
         def cond(prompt, cache, index, log_probs):
             if end_token_id is None:
                 return True
             # Stop if all sequences have produced a *new* end_token_id.
             end_tokens = (prompt == end_token_id) & (~mask)
-            prompt_done = tf.reduce_any(end_tokens, axis=-1)
-            return not tf.reduce_all(prompt_done)
+            prompt_done = ops.any(end_tokens, axis=-1)
+            return ops.logical_not(ops.all(prompt_done))
 
         def body(prompt, cache, index, log_probs):
             # Compute the softmax distribution for the next token.
             logits, _, cache = next(prompt, cache, index)
-            vocab_size = tf.shape(logits)[-1]
+            vocab_size = ops.shape(logits)[-1]
             probs = keras.activations.softmax(logits / self.temperature)
 
             # Compute the running log-likelihood of each new candidate.
-            next_log_probs = tf.math.log(probs) + log_probs[..., tf.newaxis]
+            next_log_probs = ops.log(probs) + log_probs[..., None]
             # Reshape `preds` to shape `(batch_size, num_beams * vocab_size)`.
-            next_log_probs = tf.reshape(next_log_probs, shape=[batch_size, -1])
+            next_log_probs = ops.reshape(next_log_probs, [batch_size, -1])
 
             # Compute the top beam indices and next tokens.
-            next_log_probs, indices = tf.math.top_k(
+            next_log_probs, indices = ops.top_k(
                 next_log_probs, k=self.num_beams, sorted=False
             )
             beam_indices = indices // vocab_size
             next_token = flatten_beams(indices % vocab_size)
-            # Ensure shape is `[None]`, otherwise it causes issues after
-            # converting to TFLite.
-            next_token = tf.ensure_shape(next_token, [None])
             # We need `ensure_shape` as `top_k` will change the static shape.
             next_log_probs = flatten_beams(next_log_probs)
-            log_probs = tf.ensure_shape(next_log_probs, log_probs.shape)
+            # Work around for top_k output shape on tf backend.
+            if isinstance(log_probs, tf.Tensor):
+                log_probs = tf.ensure_shape(next_log_probs, log_probs.shape)
+            else:
+                log_probs = next_log_probs
 
             def gather_beams(x):
                 x = unflatten_beams(x)
-                x = tf.gather(x, beam_indices, axis=1, batch_dims=1)
+                indices = beam_indices
+                for axis in range(2, len(x.shape)):
+                    indices = ops.expand_dims(indices, axis=axis)
+                x = ops.take_along_axis(x, indices, axis=1)
                 return flatten_beams(x)
 
             prompt = gather_beams(prompt)
-            cache = tf.nest.map_structure(gather_beams, cache)
+            if has_cache:
+                cache = tf.nest.map_structure(gather_beams, cache)
 
             # Update each beam with the next token.
-            next_token = tf.cast(next_token, prompt.dtype)
+            next_token = ops.cast(next_token, prompt.dtype)
             # Don't overwrite anywhere mask is True.
-            next_token = tf.where(mask[:, index], prompt[:, index], next_token)
+            next_token = ops.where(mask[:, index], prompt[:, index], next_token)
             # Update the prompt with the next token.
-            next_token = next_token[:, tf.newaxis]
-            prompt = dynamic_update_slice(prompt, next_token, [0, index])
+            next_token = next_token[:, None]
+            prompt = ops.slice_update(prompt, [0, index], next_token)
             # Return the iteration of the loop state.
             return (prompt, cache, index + 1, log_probs)
 
-        prompt, _, _, log_probs = tf.while_loop(
+        prompt, _, _, log_probs = self.run_loop(
             cond=cond,
             body=body,
             loop_vars=(prompt, cache, index, log_probs),
@@ -213,21 +215,23 @@ class BeamSampler(Sampler):
         all_log_probs = unflatten_beams(log_probs)
 
         if self.return_all_beams:
-            sorted_indices = tf.argsort(
-                all_log_probs, axis=-1, direction="DESCENDING"
-            )
-            sorted_log_probs = tf.gather(
-                all_log_probs, sorted_indices, axis=-1, batch_dims=1
+            sorted_indices = ops.argsort(-all_log_probs, axis=-1)
+            sorted_log_probs = ops.take_along_axis(
+                all_log_probs,
+                sorted_indices,
+                axis=1,
             )
-            sorted_prompts = tf.gather(
-                all_prompts, sorted_indices, axis=1, batch_dims=1
+            sorted_prompts = ops.take_along_axis(
+                all_prompts,
+                ops.expand_dims(sorted_indices, -1),
+                axis=1,
             )
             return sorted_prompts, sorted_log_probs
         else:
             # Gather the top beam at each batch index.
-            top_beams = tf.math.argmax(all_log_probs, axis=-1)[:, tf.newaxis]
-            prompt = tf.gather(all_prompts, top_beams, axis=1, batch_dims=1)
-            return tf.squeeze(prompt, axis=1)
+            top_beams = ops.argmax(all_log_probs, axis=-1)[:, None, None]
+            prompt = ops.take_along_axis(all_prompts, top_beams, axis=1)
+            return ops.squeeze(prompt, axis=1)
 
     def get_config(self):
         config = super().get_config()
diff --git a/keras_nlp/samplers/contrastive_sampler.py b/keras_nlp/samplers/contrastive_sampler.py
index 5a5f1fd5..ab4bd6fe 100644
--- a/keras_nlp/samplers/contrastive_sampler.py
+++ b/keras_nlp/samplers/contrastive_sampler.py
@@ -14,14 +14,13 @@
 """Contrastive Sampler."""
 
 import tensorflow as tf
-from tensorflow.compiler.tf2xla.python.xla import dynamic_update_slice
 
 from keras_nlp.api_export import keras_nlp_export
 from keras_nlp.backend import keras
+from keras_nlp.backend import ops
 from keras_nlp.samplers.sampler import Sampler
 from keras_nlp.samplers.sampler import call_args_docstring
 from keras_nlp.utils.python_utils import format_docstring
-from keras_nlp.utils.tensor_utils import assert_tf_backend
 
 
 @format_docstring(call_args=call_args_docstring)
@@ -79,10 +78,6 @@ class ContrastiveSampler(Sampler):
         seed=None,
         **kwargs,
     ):
-        # Temporarily turn off beam search in other backends.
-        # No technical blockers here, just need tf -> ops rewrite.
-        assert_tf_backend(self.__class__.__name__)
-
         super().__init__(**kwargs)
         self.k = k
         self.alpha = alpha
@@ -103,39 +98,39 @@ class ContrastiveSampler(Sampler):
                 "`ContrastiveSampler` requires passing a `hidden_states`, but"
                 "received `None`."
             )
-        batch_size, max_length = tf.shape(prompt)[0], tf.shape(prompt)[1]
-        # Make sure max length and start index are the same dtype.
-        index = tf.cast(index, max_length.dtype)
+        batch_size, max_length = ops.shape(prompt)[0], ops.shape(prompt)[1]
+        index = ops.cast(index, "int32")
 
         def create_beams(x):
             """Add initial beam state."""
-            x = tf.repeat(x, self.k, axis=0)
-            flat_shape = [batch_size * self.k] + x.shape.as_list()[1:]
-            return tf.reshape(x, shape=flat_shape)
+            x = ops.repeat(x, self.k, axis=0)
+            flat_shape = (batch_size * self.k,) + tuple(x.shape)[1:]
+            return ops.reshape(x, flat_shape)
 
         def flatten_beams(x):
             """Combine the beam dim and batch dim."""
-            flat_shape = [batch_size * self.k] + x.shape.as_list()[2:]
-            return tf.reshape(x, shape=flat_shape)
+            flat_shape = (batch_size * self.k,) + tuple(x.shape)[2:]
+            return ops.reshape(x, flat_shape)
 
         def unflatten_beams(x):
             """Separate the beam dim and batch dim."""
-            unflat_shape = [batch_size, self.k] + x.shape.as_list()[1:]
-            return tf.reshape(x, shape=unflat_shape)
+            unflat_shape = (batch_size, self.k) + tuple(x.shape)[1:]
+            return ops.reshape(x, unflat_shape)
 
-        mask = tf.zeros_like(prompt, dtype="bool") if mask is None else mask
+        mask = ops.zeros_like(prompt, dtype="bool") if mask is None else mask
         # Compute initial logits.
         logits, _, cache = next(prompt, cache, index)
-        # `tf.while_loop` will not accept `None` as a value for `loop_vars`.
-        cache = () if cache is None else cache
+        # `ops.while_loop` will not accept `None` as a value for `loop_vars`.
+        has_cache = cache is not None
+        cache = cache if has_cache else ()
 
         def cond(prompt, cache, index, logits, hidden_states):
             if end_token_id is None:
                 return True
             # Stop if all sequences have produced a *new* end_token_id.
             end_tokens = (prompt == end_token_id) & (~mask)
-            prompt_done = tf.reduce_any(end_tokens, axis=-1)
-            return not tf.reduce_all(prompt_done)
+            prompt_done = ops.any(end_tokens, axis=-1)
+            return ops.logical_not(ops.all(prompt_done))
 
         def body(prompt, cache, index, logits, hidden_states):
             # Compute the softmax distribution for the next token.
@@ -146,23 +141,25 @@ class ContrastiveSampler(Sampler):
             prompt_beams = create_beams(prompt)
             mask_beams = create_beams(mask)
             hidden_states_beams = create_beams(hidden_states)
-            cache_beams = tf.nest.map_structure(create_beams, cache)
+            cache_beams = None
+            if has_cache:
+                cache_beams = tf.nest.map_structure(create_beams, cache)
 
             # Get top-k candidate tokens and their probabilities.
-            top_k_probabilities, top_k_indices = tf.math.top_k(
+            top_k_probabilities, top_k_indices = ops.top_k(
                 probabilities, k=self.k, sorted=False
             )
             next_token_probabilities = flatten_beams(top_k_probabilities)
             next_token = flatten_beams(top_k_indices)
-            next_token = tf.cast(next_token, prompt.dtype)
-            next_token = tf.where(
+            next_token = ops.cast(next_token, prompt.dtype)
+            next_token = ops.where(
                 mask_beams[:, index], prompt_beams[:, index], next_token
             )
 
             # Update the prompt with the next token.
-            next_token = next_token[:, tf.newaxis]
-            prompt_beams = dynamic_update_slice(
-                prompt_beams, next_token, [0, index]
+            next_token = ops.expand_dims(next_token, -1)
+            prompt_beams = ops.slice_update(
+                prompt_beams, [0, index], next_token
             )
 
             # Compute the logits and hidden states for top-k candidate tokens.
@@ -175,14 +172,13 @@ class ContrastiveSampler(Sampler):
             similarity_scores = self.similarity(
                 hidden_states_beams, next_hidden_states_beams
             )
-            max_similarity_scores = tf.cast(
-                tf.reduce_max(similarity_scores[:, :index], axis=1),
+            # Replace all future indices with -1, the lowest similarity score.
+            score_mask = ops.expand_dims(ops.arange(max_length) < index, 0)
+            similarity_scores = ops.where(score_mask, similarity_scores, -1)
+            max_similarity_scores = ops.cast(
+                ops.max(similarity_scores, axis=1),
                 dtype=next_token_probabilities.dtype,
             )
-            if index == 0:
-                # If the index is 0, there is no previous states so we set
-                # `max_similarity_scores` the same for all beams.
-                max_similarity_scores = tf.zeros_like(max_similarity_scores)
             # The final score of each candidate token is weighted sum of
             # probability and similarity against previous tokens.
             accumulated_scores = (
@@ -197,16 +193,18 @@ class ContrastiveSampler(Sampler):
             unflat_next_hidden_states = unflatten_beams(
                 next_hidden_states_beams
             )
-            unflat_cache = tf.nest.map_structure(unflatten_beams, cache_beams)
-            best_token_indices = tf.math.argmax(unflat_score, axis=1)
+            best_token_indices = ops.argmax(unflat_score, axis=1)
 
             def gather_best_token(beams):
-                return tf.gather(
+                indices = best_token_indices
+                for axis in range(1, len(beams.shape)):
+                    indices = ops.expand_dims(indices, axis=axis)
+                best = ops.take_along_axis(
                     beams,
-                    best_token_indices,
+                    indices,
                     axis=1,
-                    batch_dims=1,
                 )
+                return ops.squeeze(best, axis=1)
 
             prompt = gather_best_token(unflat_prompt)
             # We avoid recomputing forward pass for each token by updating the
@@ -214,16 +212,18 @@ class ContrastiveSampler(Sampler):
             # next iteration step.
             logits = gather_best_token(unflat_next_logits)
             next_hidden_states = gather_best_token(unflat_next_hidden_states)
-            cache = tf.nest.map_structure(gather_best_token, unflat_cache)
+            if has_cache:
+                cache = tf.nest.map_structure(unflatten_beams, cache_beams)
+                cache = tf.nest.map_structure(gather_best_token, cache)
 
-            hidden_states = dynamic_update_slice(
+            hidden_states = ops.slice_update(
                 hidden_states,
-                next_hidden_states[:, tf.newaxis, :],
                 [0, index, 0],
+                next_hidden_states[:, None, :],
             )
             return (prompt, cache, index + 1, logits, hidden_states)
 
-        prompt, _, _, _, _ = tf.while_loop(
+        prompt, _, _, _, _ = self.run_loop(
             cond=cond,
             body=body,
             loop_vars=(prompt, cache, index, logits, hidden_states),
@@ -232,10 +232,10 @@ class ContrastiveSampler(Sampler):
         return prompt
 
     def similarity(self, h1, h2):
-        h2 = h2[..., tf.newaxis]
-        return tf.squeeze(tf.matmul(h1, h2), axis=-1) / (
-            tf.norm(h1, axis=-1) * tf.norm(h2, axis=-2)
-        )
+        h2 = ops.expand_dims(h2, -1)
+        h1_norm = ops.sqrt(ops.sum(h1 * h1, axis=-1))
+        h2_norm = ops.sqrt(ops.sum(h2 * h2, axis=-2))
+        return ops.squeeze(ops.matmul(h1, h2), axis=-1) / (h1_norm * h2_norm)
 
     def get_config(self):
         config = super().get_config()


Test patch:
diff --git a/keras_nlp/layers/modeling/masked_lm_head_test.py b/keras_nlp/layers/modeling/masked_lm_head_test.py
index f5c3b9d0..7d85da5b 100644
--- a/keras_nlp/layers/modeling/masked_lm_head_test.py
+++ b/keras_nlp/layers/modeling/masked_lm_head_test.py
@@ -29,7 +29,7 @@ class MaskedLMHeadTest(TestCase):
         )
         encoded_tokens = keras.Input(shape=(10, 16))
         positions = keras.Input(shape=(5,), dtype="int32")
-        outputs = head(encoded_tokens, masked_positions=positions)
+        outputs = head(encoded_tokens, mask_positions=positions)
         model = keras.Model((encoded_tokens, positions), outputs)
 
         token_data = ops.random.uniform(shape=(4, 10, 16))
@@ -48,7 +48,7 @@ class MaskedLMHeadTest(TestCase):
         # need to support this in the layer.
         sequence = keras.Input(shape=(10, 32))
         positions = keras.Input(shape=(5,), dtype="int32")
-        outputs = head(sequence, masked_positions=positions)
+        outputs = head(sequence, mask_positions=positions)
         model = keras.Model((sequence, positions), outputs)
         sequence_data = ops.random.uniform(shape=(4, 10, 32))
         position_data = ops.random.randint(minval=0, maxval=10, shape=(4, 5))
@@ -106,7 +106,7 @@ class MaskedLMHeadTest(TestCase):
         )
         encoded_tokens = keras.Input(shape=(10, 16))
         positions = keras.Input(shape=(5,), dtype="int32")
-        outputs = head(encoded_tokens, masked_positions=positions)
+        outputs = head(encoded_tokens, mask_positions=positions)
         model = keras.Model((encoded_tokens, positions), outputs)
 
         token_data = ops.random.uniform(shape=(4, 10, 16))
@@ -126,7 +126,7 @@ class MaskedLMHeadTest(TestCase):
         )
         encoded_tokens = keras.Input(shape=(10, 16))
         positions = keras.Input(shape=(5,), dtype="int32")
-        outputs = head(encoded_tokens, masked_positions=positions)
+        outputs = head(encoded_tokens, mask_positions=positions)
         model = keras.Model((encoded_tokens, positions), outputs)
 
         token_data = ops.random.uniform(shape=(4, 10, 16))
diff --git a/keras_nlp/metrics/bleu_test.py b/keras_nlp/metrics/bleu_test.py
index 756c0688..b1d8dc49 100644
--- a/keras_nlp/metrics/bleu_test.py
+++ b/keras_nlp/metrics/bleu_test.py
@@ -22,7 +22,6 @@ from keras_nlp.tests.test_case import TestCase
 from keras_nlp.tokenizers.byte_tokenizer import ByteTokenizer
 
 
-@pytest.mark.tf_only
 class BleuTest(TestCase):
     def test_initialization(self):
         bleu = Bleu()
@@ -69,64 +68,38 @@ class BleuTest(TestCase):
         bleu_val = bleu(y_true, y_pred)
         self.assertAlmostEqual(bleu_val, 0.243, delta=1e-3)
 
-    def test_1d_tensor_input(self):
-        bleu = Bleu()
-        y_true = tf.ragged.constant(
-            [
-                ["He eats a sweet apple."],
-                ["Silicon Valley is one of my favourite shows!"],
-            ]
-        )
-        y_pred = tf.constant(
-            [
-                "He He He eats sweet apple which is a fruit.",
-                "I love Silicon Valley, it's one of my favourite shows.",
-            ]
-        )
-
-        bleu_val = bleu(y_true, y_pred)
-        self.assertAlmostEqual(bleu_val, 0.243, delta=1e-3)
-
-    def test_2d_tensor_input(self):
-        bleu = Bleu()
-        y_true = tf.constant(
-            [
-                [["He eats a sweet apple."]],
-                [["Silicon Valley is one of my favourite shows!"]],
-            ]
-        )
-        y_pred = tf.constant(
-            [
-                ["He He He eats sweet apple which is a fruit."],
-                ["I love Silicon Valley, it's one of my favourite shows."],
-            ]
-        )
-
-        bleu_val = bleu(y_true, y_pred)
-        self.assertAlmostEqual(bleu_val, 0.243, delta=1e-3)
-
     def test_custom_tokenizer(self):
         byte_tokenizer = ByteTokenizer()
         bleu = Bleu(tokenizer=byte_tokenizer)
-        y_true = tf.ragged.constant(
-            [
-                ["He eats a sweet apple."],
-                ["Silicon Valley is one of my favourite shows!"],
-            ]
-        )
-        y_pred = tf.constant(
-            [
-                "He He He eats sweet apple which is a fruit.",
-                "I love Silicon Valley, it's one of my favourite shows.",
-            ]
-        )
+        y_true = [
+            ["He eats a sweet apple."],
+            ["Silicon Valley is one of my favourite shows!"],
+        ]
+        y_pred = [
+            "He He He eats sweet apple which is a fruit.",
+            "I love Silicon Valley, it's one of my favourite shows.",
+        ]
 
         bleu_val = bleu(y_true, y_pred)
         self.assertAlmostEqual(bleu_val, 0.609, delta=1e-3)
 
     def test_different_order(self):
         bleu = Bleu(max_order=5)
-        y_true = tf.ragged.constant(
+        y_true = [
+            ["He eats a sweet apple."],
+            ["Silicon Valley is one of my favourite shows!"],
+        ]
+        y_pred = [
+            "He He He eats sweet apple which is a fruit.",
+            "I love Silicon Valley, it's one of my favourite shows.",
+        ]
+
+        bleu_val = bleu(y_true, y_pred)
+        self.assertAlmostEqual(bleu_val, 0.188, delta=1e-3)
+
+    def test_tensor_input(self):
+        bleu = Bleu()
+        y_true = tf.constant(
             [
                 ["He eats a sweet apple."],
                 ["Silicon Valley is one of my favourite shows!"],
@@ -140,8 +113,9 @@ class BleuTest(TestCase):
         )
 
         bleu_val = bleu(y_true, y_pred)
-        self.assertAlmostEqual(bleu_val, 0.188, delta=1e-3)
+        self.assertAlmostEqual(bleu_val, 0.243, delta=1e-3)
 
+    @pytest.mark.tf_only  # string model output only applies to tf.
     def test_model_compile(self):
         inputs = keras.Input(shape=(), dtype="string")
         outputs = keras.layers.Identity()(inputs)
@@ -166,18 +140,14 @@ class BleuTest(TestCase):
 
     def test_reset_state(self):
         bleu = Bleu()
-        y_true = tf.ragged.constant(
-            [
-                ["He eats a sweet apple."],
-                ["Silicon Valley is one of my favourite shows!"],
-            ]
-        )
-        y_pred = tf.constant(
-            [
-                "He He He eats sweet apple which is a fruit.",
-                "I love Silicon Valley, it's one of my favourite shows.",
-            ]
-        )
+        y_true = [
+            ["He eats a sweet apple."],
+            ["Silicon Valley is one of my favourite shows!"],
+        ]
+        y_pred = [
+            "He He He eats sweet apple which is a fruit.",
+            "I love Silicon Valley, it's one of my favourite shows.",
+        ]
 
         bleu.update_state(y_true, y_pred)
         bleu_val = bleu.result()
@@ -189,25 +159,21 @@ class BleuTest(TestCase):
 
     def test_update_state(self):
         bleu = Bleu()
-        y_true_1 = tf.ragged.constant(
-            [
-                ["He eats a sweet apple."],
-                ["Silicon Valley is one of my favourite shows!"],
-            ]
-        )
-        y_pred_1 = tf.constant(
-            [
-                "He He He eats sweet apple which is a fruit.",
-                "I love Silicon Valley, it's one of my favourite shows.",
-            ]
-        )
+        y_true_1 = [
+            ["He eats a sweet apple."],
+            ["Silicon Valley is one of my favourite shows!"],
+        ]
+        y_pred_1 = [
+            "He He He eats sweet apple which is a fruit.",
+            "I love Silicon Valley, it's one of my favourite shows.",
+        ]
 
         bleu.update_state(y_true_1, y_pred_1)
         bleu_val = bleu.result()
         self.assertAlmostEqual(bleu_val, 0.243, delta=1e-3)
 
-        y_true_2 = tf.constant(["Virat Kohli is the GOAT."])
-        y_pred_2 = tf.constant("Virat Kohli is the greatest of all time!")
+        y_true_2 = ["Virat Kohli is the GOAT."]
+        y_pred_2 = "Virat Kohli is the greatest of all time!"
 
         bleu.update_state(y_true_2, y_pred_2)
         bleu_val = bleu.result()
diff --git a/keras_nlp/metrics/edit_distance_test.py b/keras_nlp/metrics/edit_distance_test.py
index 620d6519..c463fc55 100644
--- a/keras_nlp/metrics/edit_distance_test.py
+++ b/keras_nlp/metrics/edit_distance_test.py
@@ -21,7 +21,6 @@ from keras_nlp.metrics.edit_distance import EditDistance
 from keras_nlp.tests.test_case import TestCase
 
 
-@pytest.mark.tf_only
 class EditDistanceTest(TestCase):
     def test_initialization(self):
         edit_distance = EditDistance()
@@ -51,46 +50,30 @@ class EditDistanceTest(TestCase):
         edit_distance_val = edit_distance(y_true, y_pred)
         self.assertAlmostEqual(edit_distance_val, 0.733, delta=1e-3)
 
-    def test_rank_1_tensor_input_normalize(self):
-        edit_distance = EditDistance()
-        y_true = tf.strings.split(
-            "the tiny little cat was found under the big funny bed"
-        )
-        y_pred = tf.strings.split("the cat was found under the bed")
-
-        edit_distance_val = edit_distance(y_true, y_pred)
-        self.assertAlmostEqual(edit_distance_val, 0.364, delta=1e-3)
-
-    def test_rank_2_tensor_input_normalize(self):
-        edit_distance = EditDistance()
-        y_true = tf.strings.split(
-            [
-                "the tiny little cat was found under the big funny bed",
-                "it is sunny today",
-            ]
-        )
-        y_pred = tf.strings.split(
-            [
-                "the cat was found under the bed",
-                "it is sunny but with a hint of cloud cover",
-            ]
-        )
+    def test_1d_list_input_normalize_false(self):
+        edit_distance = EditDistance(normalize=False)
+        y_true = "the tiny little cat was found under the big funny bed".split()
+        y_pred = "the cat was found under the bed".split()
 
         edit_distance_val = edit_distance(y_true, y_pred)
-        self.assertAlmostEqual(edit_distance_val, 0.733, delta=1e-3)
+        self.assertAlmostEqual(edit_distance_val, 4.0, delta=1e-3)
 
-    def test_rank_1_tensor_input_normalize_false(self):
+    def test_2d_list_input_normalize_false(self):
         edit_distance = EditDistance(normalize=False)
-        y_true = tf.strings.split(
-            "the tiny little cat was found under the big funny bed"
-        )
-        y_pred = tf.strings.split("the cat was found under the bed")
+        y_true = [
+            "the tiny little cat was found under the big funny bed".split(),
+            "it is sunny today".split(),
+        ]
+        y_pred = [
+            "the cat was found under the bed".split(),
+            "it is sunny but with a hint of cloud cover".split(),
+        ]
 
         edit_distance_val = edit_distance(y_true, y_pred)
-        self.assertAlmostEqual(edit_distance_val, 4.0, delta=1e-3)
+        self.assertAlmostEqual(edit_distance_val, 5.5, delta=1e-3)
 
-    def test_rank_2_tensor_input_normalize_false(self):
-        edit_distance = EditDistance(normalize=False)
+    def test_tensor_input(self):
+        edit_distance = EditDistance()
         y_true = tf.strings.split(
             [
                 "the tiny little cat was found under the big funny bed",
@@ -105,8 +88,9 @@ class EditDistanceTest(TestCase):
         )
 
         edit_distance_val = edit_distance(y_true, y_pred)
-        self.assertAlmostEqual(edit_distance_val, 5.5, delta=1e-3)
+        self.assertAlmostEqual(edit_distance_val, 0.733, delta=1e-3)
 
+    @pytest.mark.tf_only  # string model output only applies to tf.
     def test_model_compile_normalize(self):
         inputs = keras.Input(shape=(None,), dtype="string")
         outputs = keras.layers.Identity()(inputs)
@@ -122,6 +106,7 @@ class EditDistanceTest(TestCase):
         output = model.compute_metrics(x, y, y_pred, sample_weight=None)
         self.assertAlmostEqual(output["edit_distance"], 0.364, delta=1e-3)
 
+    @pytest.mark.tf_only  # string model output only applies to tf.
     def test_model_compile_normalize_false(self):
         inputs = keras.Input(shape=(None,), dtype="string")
         outputs = keras.layers.Identity()(inputs)
@@ -137,20 +122,26 @@ class EditDistanceTest(TestCase):
         output = model.compute_metrics(x, y, y_pred, sample_weight=None)
         self.assertAlmostEqual(output["edit_distance"], 4.0, delta=1e-3)
 
-    def test_reset_state_normalize(self):
+    def test_rank_1_tensor_input_normalize(self):
         edit_distance = EditDistance()
         y_true = tf.strings.split(
-            [
-                "the tiny little cat was found under the big funny bed",
-                "it is sunny today",
-            ]
-        )
-        y_pred = tf.strings.split(
-            [
-                "the cat was found under the bed",
-                "it is sunny but with a hint of cloud cover",
-            ]
+            "the tiny little cat was found under the big funny bed"
         )
+        y_pred = tf.strings.split("the cat was found under the bed")
+
+        edit_distance_val = edit_distance(y_true, y_pred)
+        self.assertAlmostEqual(edit_distance_val, 0.364, delta=1e-3)
+
+    def test_reset_state_normalize(self):
+        edit_distance = EditDistance()
+        y_true = [
+            "the tiny little cat was found under the big funny bed".split(),
+            "it is sunny today".split(),
+        ]
+        y_pred = [
+            "the cat was found under the bed".split(),
+            "it is sunny but with a hint of cloud cover".split(),
+        ]
 
         edit_distance.update_state(y_true, y_pred)
         edit_distance_val = edit_distance.result()
@@ -162,18 +153,14 @@ class EditDistanceTest(TestCase):
 
     def test_update_state_normalize(self):
         edit_distance = EditDistance()
-        y_true_1 = tf.strings.split(
-            [
-                "the tiny little cat was found under the big funny bed",
-                "it is sunny today",
-            ]
-        )
-        y_pred_1 = tf.strings.split(
-            [
-                "the cat was found under the bed",
-                "it is sunny but with a hint of cloud cover",
-            ]
-        )
+        y_true_1 = [
+            "the tiny little cat was found under the big funny bed".split(),
+            "it is sunny today".split(),
+        ]
+        y_pred_1 = [
+            "the cat was found under the bed".split(),
+            "it is sunny but with a hint of cloud cover".split(),
+        ]
 
         edit_distance.update_state(y_true_1, y_pred_1)
         edit_distance_val = edit_distance.result()
@@ -188,18 +175,14 @@ class EditDistanceTest(TestCase):
 
     def test_update_state_normalize_false(self):
         edit_distance = EditDistance(normalize=False)
-        y_true_1 = tf.strings.split(
-            [
-                "the tiny little cat was found under the big funny bed",
-                "it is sunny today",
-            ]
-        )
-        y_pred_1 = tf.strings.split(
-            [
-                "the cat was found under the bed",
-                "it is sunny but with a hint of cloud cover",
-            ]
-        )
+        y_true_1 = [
+            "the tiny little cat was found under the big funny bed".split(),
+            "it is sunny today".split(),
+        ]
+        y_pred_1 = [
+            "the cat was found under the bed".split(),
+            "it is sunny but with a hint of cloud cover".split(),
+        ]
 
         edit_distance.update_state(y_true_1, y_pred_1)
         edit_distance_val = edit_distance.result()
diff --git a/keras_nlp/metrics/perplexity_test.py b/keras_nlp/metrics/perplexity_test.py
index aae1e09d..cdf69c9b 100644
--- a/keras_nlp/metrics/perplexity_test.py
+++ b/keras_nlp/metrics/perplexity_test.py
@@ -13,14 +13,11 @@
 # limitations under the License.
 
 """Tests for Perplexity."""
-import pytest
-import tensorflow as tf
-
+from keras_nlp.backend import ops
 from keras_nlp.metrics.perplexity import Perplexity
 from keras_nlp.tests.test_case import TestCase
 
 
-@pytest.mark.tf_only
 class PerplexityTest(TestCase):
     def test_vars_after_initializing_class(self):
         perplexity = Perplexity()
@@ -28,8 +25,8 @@ class PerplexityTest(TestCase):
 
     def test_from_logits_without_masking(self):
         perplexity = Perplexity(from_logits=True)
-        y_true = tf.constant([[1, 3, 0], [2, 1, 3]])
-        y_pred = tf.constant(
+        y_true = ops.array([[1, 3, 0], [2, 1, 3]])
+        y_pred = ops.array(
             [
                 [
                     [1.034, 4.797, 2.82, 1.154],
@@ -50,8 +47,8 @@ class PerplexityTest(TestCase):
     def test_from_logits_with_sample_weight(self):
         perplexity = Perplexity(from_logits=True)
 
-        y_true = tf.constant([[1, 3, 0], [2, 1, 3]])
-        y_pred = tf.constant(
+        y_true = ops.array([[1, 3, 0], [2, 1, 3]])
+        y_pred = ops.array(
             [
                 [
                     [1.034, 4.797, 2.82, 1.154],
@@ -65,7 +62,7 @@ class PerplexityTest(TestCase):
                 ],
             ]
         )
-        sample_wt = tf.cast(y_true != 0, "int32")
+        sample_wt = ops.cast(y_true != 0, "int32")
 
         perplexity_val = perplexity(y_true, y_pred, sample_wt)
         self.assertAlmostEqual(perplexity_val, 2.8789, delta=1e-3)
@@ -73,8 +70,8 @@ class PerplexityTest(TestCase):
     def test_from_logits_with_mask_token_id(self):
         perplexity = Perplexity(from_logits=True, mask_token_id=0)
 
-        y_true = tf.constant([[1, 3, 0], [2, 1, 3]])
-        y_pred = tf.constant(
+        y_true = ops.array([[1, 3, 0], [2, 1, 3]])
+        y_pred = ops.array(
             [
                 [
                     [1.034, 4.797, 2.82, 1.154],
@@ -95,8 +92,8 @@ class PerplexityTest(TestCase):
     def test_from_logits_with_mask_token_id_and_sample_weight(self):
         perplexity = Perplexity(from_logits=True, mask_token_id=0)
 
-        y_true = tf.constant([[1, 3, 0], [2, 1, 3]])
-        y_pred = tf.constant(
+        y_true = ops.array([[1, 3, 0], [2, 1, 3]])
+        y_pred = ops.array(
             [
                 [
                     [1.034, 4.797, 2.82, 1.154],
@@ -110,7 +107,7 @@ class PerplexityTest(TestCase):
                 ],
             ]
         )
-        sample_weight = tf.constant([[0.5, 0.1, 0.9], [1, 0.7, 0.5]])
+        sample_weight = ops.array([[0.5, 0.1, 0.9], [1, 0.7, 0.5]])
 
         perplexity_val = perplexity(y_true, y_pred, sample_weight)
         self.assertAlmostEqual(perplexity_val, 2.9442, delta=1e-3)
@@ -118,8 +115,8 @@ class PerplexityTest(TestCase):
     def test_two_inputs_from_logits(self):
         perplexity = Perplexity(from_logits=True, mask_token_id=0)
 
-        y_true_1 = tf.constant([[1, 3, 0], [2, 1, 3]])
-        y_pred_1 = tf.constant(
+        y_true_1 = ops.array([[1, 3, 0], [2, 1, 3]])
+        y_pred_1 = ops.array(
             [
                 [
                     [1.034, 4.797, 2.82, 1.154],
@@ -137,8 +134,8 @@ class PerplexityTest(TestCase):
         perplexity_val = perplexity(y_true_1, y_pred_1)
         self.assertAlmostEqual(perplexity_val, 2.8789, delta=1e-3)
 
-        y_true_2 = tf.constant([[2, 0, 0], [1, 2, 3]])
-        y_pred_2 = tf.constant(
+        y_true_2 = ops.array([[2, 0, 0], [1, 2, 3]])
+        y_pred_2 = ops.array(
             [
                 [
                     [2.887, 0.885, 2.973, 2.582],
@@ -158,8 +155,8 @@ class PerplexityTest(TestCase):
     def test_from_probs_with_sample_weight(self):
         perplexity = Perplexity(from_logits=False)
 
-        y_true = tf.constant([[1, 3, 0], [2, 1, 3]])
-        y_pred = tf.constant(
+        y_true = ops.array([[1, 3, 0], [2, 1, 3]])
+        y_pred = ops.array(
             [
                 [
                     [1.034, 4.797, 2.82, 1.154],
@@ -173,9 +170,9 @@ class PerplexityTest(TestCase):
                 ],
             ]
         )
-        y_prob = tf.nn.softmax(y_pred, axis=-1)
+        y_prob = ops.softmax(y_pred, axis=-1)
 
-        sample_wt = tf.cast(y_true != 0, "int32")
+        sample_wt = ops.cast(y_true != 0, "int32")
 
         perplexity_val = perplexity(y_true, y_prob, sample_wt)
         self.assertAlmostEqual(perplexity_val, 2.8789, delta=1e-3)
@@ -183,8 +180,8 @@ class PerplexityTest(TestCase):
     def test_from_probs_with_pad_token(self):
         perplexity = Perplexity(from_logits=False, mask_token_id=0)
 
-        y_true = tf.constant([[1, 3, 0], [2, 1, 3]])
-        y_pred = tf.constant(
+        y_true = ops.array([[1, 3, 0], [2, 1, 3]])
+        y_pred = ops.array(
             [
                 [
                     [1.034, 4.797, 2.82, 1.154],
@@ -198,14 +195,14 @@ class PerplexityTest(TestCase):
                 ],
             ]
         )
-        y_prob = tf.nn.softmax(y_pred, axis=-1)
+        y_prob = ops.softmax(y_pred, axis=-1)
 
         perplexity_val = perplexity(y_true, y_prob)
         self.assertAlmostEqual(perplexity_val, 2.8789, delta=1e-3)
 
     def test_reset_state(self):
-        y_true = tf.constant([[1, 3, 0], [2, 1, 3]])
-        y_pred = tf.constant(
+        y_true = ops.array([[1, 3, 0], [2, 1, 3]])
+        y_pred = ops.array(
             [
                 [
                     [1.034, 4.797, 2.82, 1.154],
@@ -231,8 +228,8 @@ class PerplexityTest(TestCase):
     def test_update_state(self):
         perplexity = Perplexity(from_logits=True, mask_token_id=0)
 
-        y_true_1 = tf.constant([[1, 3, 0], [2, 1, 3]])
-        y_pred_1 = tf.constant(
+        y_true_1 = ops.array([[1, 3, 0], [2, 1, 3]])
+        y_pred_1 = ops.array(
             [
                 [
                     [1.034, 4.797, 2.82, 1.154],
@@ -251,8 +248,8 @@ class PerplexityTest(TestCase):
         perplexity_val = perplexity.result()
         self.assertAlmostEqual(perplexity_val, 2.8789, delta=1e-3)
 
-        y_true_2 = tf.constant([[2, 0, 0], [1, 2, 3]])
-        y_pred_2 = tf.constant(
+        y_true_2 = ops.array([[2, 0, 0], [1, 2, 3]])
+        y_pred_2 = ops.array(
             [
                 [
                     [2.887, 0.885, 2.973, 2.582],
diff --git a/keras_nlp/metrics/rouge_l_test.py b/keras_nlp/metrics/rouge_l_test.py
index 7b0818d3..4b1206cb 100644
--- a/keras_nlp/metrics/rouge_l_test.py
+++ b/keras_nlp/metrics/rouge_l_test.py
@@ -13,14 +13,12 @@
 # limitations under the License.
 
 """Tests for RougeL."""
-import pytest
 import tensorflow as tf
 
 from keras_nlp.metrics.rouge_l import RougeL
 from keras_nlp.tests.test_case import TestCase
 
 
-@pytest.mark.tf_only
 class RougeLTest(TestCase):
     def test_initialization(self):
         rouge = RougeL()
@@ -77,35 +75,13 @@ class RougeLTest(TestCase):
             {"precision": 1.0, "recall": 0.689393, "f1_score": 0.807486},
         )
 
-    def test_rank_2_input(self):
-        rouge = RougeL(use_stemmer=False)
-        y_true = tf.constant(
-            [
-                ["the tiny little cat was found under the big funny bed"],
-                ["i really love contributing to KerasNLP"],
-            ]
-        )
-        y_pred = tf.constant(
-            [["the cat was under the bed"], ["i love contributing to KerasNLP"]]
-        )
-
-        rouge_val = rouge(y_true, y_pred)
-        self.assertAllClose(
-            rouge_val,
-            {"precision": 1.0, "recall": 0.689393, "f1_score": 0.807486},
-        )
-
     def test_reset_state(self):
         rouge = RougeL()
-        y_true = tf.constant(
-            ["hey, this is great fun", "i love contributing to KerasNLP"]
-        )
-        y_pred = tf.constant(
-            [
-                "great fun indeed",
-                "KerasNLP is awesome, i love contributing to it",
-            ]
-        )
+        y_true = ["hey, this is great fun", "i love contributing to KerasNLP"]
+        y_pred = [
+            "great fun indeed",
+            "KerasNLP is awesome, i love contributing to it",
+        ]
 
         rouge.update_state(y_true, y_pred)
         rouge_val = rouge.result()
@@ -123,15 +99,14 @@ class RougeLTest(TestCase):
 
     def test_update_state(self):
         rouge = RougeL()
-        y_true_1 = tf.constant(
-            [
-                "the tiny little cat was found under the big funny bed",
-                "i really love contributing to KerasNLP",
-            ]
-        )
-        y_pred_1 = tf.constant(
-            ["the cat was under the bed", "i love contributing to KerasNLP"]
-        )
+        y_true_1 = [
+            "the tiny little cat was found under the big funny bed",
+            "i really love contributing to KerasNLP",
+        ]
+        y_pred_1 = [
+            "the cat was under the bed",
+            "i love contributing to KerasNLP",
+        ]
 
         rouge.update_state(y_true_1, y_pred_1)
         rouge_val = rouge.result()
@@ -140,8 +115,8 @@ class RougeLTest(TestCase):
             {"precision": 1.0, "recall": 0.689393, "f1_score": 0.807486},
         )
 
-        y_true_2 = tf.constant(["what is your favourite show"])
-        y_pred_2 = tf.constant(["my favourite show is silicon valley"])
+        y_true_2 = ["what is your favourite show"]
+        y_pred_2 = ["my favourite show is silicon valley"]
 
         rouge.update_state(y_true_2, y_pred_2)
         rouge_val = rouge.result()
diff --git a/keras_nlp/metrics/rouge_n_test.py b/keras_nlp/metrics/rouge_n_test.py
index ac64906a..21effb0b 100644
--- a/keras_nlp/metrics/rouge_n_test.py
+++ b/keras_nlp/metrics/rouge_n_test.py
@@ -21,7 +21,6 @@ from keras_nlp.metrics.rouge_n import RougeN
 from keras_nlp.tests.test_case import TestCase
 
 
-@pytest.mark.tf_only
 class RougeNTest(TestCase):
     def test_initialization(self):
         rouge = RougeN()
@@ -78,24 +77,7 @@ class RougeNTest(TestCase):
             {"precision": 0.575, "recall": 0.4, "f1_score": 0.466666},
         )
 
-    def test_rank_2_input(self):
-        rouge = RougeN(order=2, use_stemmer=False)
-        y_true = tf.constant(
-            [
-                ["the tiny little cat was found under the big funny bed"],
-                ["i really love contributing to KerasNLP"],
-            ]
-        )
-        y_pred = tf.constant(
-            [["the cat was under the bed"], ["i love contributing to KerasNLP"]]
-        )
-
-        rouge_val = rouge(y_true, y_pred)
-        self.assertAllClose(
-            rouge_val,
-            {"precision": 0.575, "recall": 0.4, "f1_score": 0.466666},
-        )
-
+    @pytest.mark.tf_only  # string model output only applies to tf.
     def test_model_compile(self):
         inputs = keras.Input(shape=(None,), dtype="string")
         outputs = keras.layers.Identity()(inputs)
@@ -117,15 +99,14 @@ class RougeNTest(TestCase):
 
     def test_different_order(self):
         rouge = RougeN(order=3, use_stemmer=False)
-        y_true = tf.constant(
-            [
-                "the tiny little cat was found under the big funny bed",
-                "i really love contributing to KerasNLP",
-            ]
-        )
-        y_pred = tf.constant(
-            ["the cat was under the bed", "i love contributing to KerasNLP"]
-        )
+        y_true = [
+            "the tiny little cat was found under the big funny bed",
+            "i really love contributing to KerasNLP",
+        ]
+        y_pred = [
+            "the cat was under the bed",
+            "i love contributing to KerasNLP",
+        ]
 
         rouge_val = rouge(y_true, y_pred)
         self.assertAllClose(
@@ -135,15 +116,11 @@ class RougeNTest(TestCase):
 
     def test_reset_state(self):
         rouge = RougeN()
-        y_true = tf.constant(
-            ["hey, this is great fun", "i love contributing to KerasNLP"]
-        )
-        y_pred = tf.constant(
-            [
-                "great fun indeed",
-                "KerasNLP is awesome, i love contributing to it",
-            ]
-        )
+        y_true = ["hey, this is great fun", "i love contributing to KerasNLP"]
+        y_pred = [
+            "great fun indeed",
+            "KerasNLP is awesome, i love contributing to it",
+        ]
 
         rouge.update_state(y_true, y_pred)
         rouge_val = rouge.result()
@@ -161,15 +138,14 @@ class RougeNTest(TestCase):
 
     def test_update_state(self):
         rouge = RougeN()
-        y_true_1 = tf.constant(
-            [
-                "the tiny little cat was found under the big funny bed",
-                "i really love contributing to KerasNLP",
-            ]
-        )
-        y_pred_1 = tf.constant(
-            ["the cat was under the bed", "i love contributing to KerasNLP"]
-        )
+        y_true_1 = [
+            "the tiny little cat was found under the big funny bed",
+            "i really love contributing to KerasNLP",
+        ]
+        y_pred_1 = [
+            "the cat was under the bed",
+            "i love contributing to KerasNLP",
+        ]
 
         rouge.update_state(y_true_1, y_pred_1)
         rouge_val = rouge.result()
@@ -178,8 +154,8 @@ class RougeNTest(TestCase):
             {"precision": 0.575, "recall": 0.4, "f1_score": 0.466666},
         )
 
-        y_true_2 = tf.constant(["what is your favourite show"])
-        y_pred_2 = tf.constant(["my favourite show is silicon valley"])
+        y_true_2 = ["what is your favourite show"]
+        y_pred_2 = ["my favourite show is silicon valley"]
 
         rouge.update_state(y_true_2, y_pred_2)
         rouge_val = rouge.result()
diff --git a/keras_nlp/samplers/beam_sampler_test.py b/keras_nlp/samplers/beam_sampler_test.py
index 88b68d87..737fe8aa 100644
--- a/keras_nlp/samplers/beam_sampler_test.py
+++ b/keras_nlp/samplers/beam_sampler_test.py
@@ -22,7 +22,6 @@ from keras_nlp.samplers.beam_sampler import BeamSampler
 from keras_nlp.tests.test_case import TestCase
 
 
-@pytest.mark.tf_only
 class BeamSamplerTest(TestCase):
     def setUp(self):
         super().setUp()
diff --git a/keras_nlp/samplers/contrastive_sampler_test.py b/keras_nlp/samplers/contrastive_sampler_test.py
index f89c708f..51d1d538 100644
--- a/keras_nlp/samplers/contrastive_sampler_test.py
+++ b/keras_nlp/samplers/contrastive_sampler_test.py
@@ -22,7 +22,6 @@ from keras_nlp.samplers.contrastive_sampler import ContrastiveSampler
 from keras_nlp.tests.test_case import TestCase
 
 
-@pytest.mark.tf_only
 class ContrastiveSamplerTest(TestCase):
     def setUp(self):
         super().setUp()