def METHOD_NAME(self, strategy):
    if (
        backend.is_tpu_strategy(strategy)
        and not tf_test_utils.is_mlir_bridge_enabled()
    ):
        self.skipTest("TPU tests require MLIR bridge")
    vocab_data = [
        [
            "earth",
            "earth",
            "earth",
            "earth",
            "wind",
            "wind",
            "wind",
            "and",
            "and",
            "fire",
        ]
    ]
    vocab_dataset = tf.data.Dataset.from_tensors(vocab_data)
    input_array = np.array(
        [
            ["earth", "wind", "and", "fire"],
            ["fire", "and", "earth", "michigan"],
        ]
    )
    input_dataset = tf.data.Dataset.from_tensor_slices(input_array).batch(
        2, drop_remainder=True
    )
    expected_output = [[2, 3, 4, 5], [5, 4, 2, 1]]
    tf.config.set_soft_device_placement(True)
    with strategy.scope():
        input_data = keras.Input(shape=(None,), dtype=tf.string)
        layer = index_lookup.IndexLookup(
            max_tokens=None,
            num_oov_indices=1,
            mask_token="",
            oov_token="[OOV]",
            vocabulary_dtype=tf.string,
        )
        layer.adapt(vocab_dataset)
        int_data = layer(input_data)
        model = keras.Model(inputs=input_data, outputs=int_data)
    model.compile(loss="mse")
    output_dataset = model.predict(input_dataset)
    self.assertAllEqual(expected_output, output_dataset)