tf-keras-nightly 2.19.0.dev2024121210__py3-none-any.whl → 2.21.0.dev2025123010__py3-none-any.whl

tf_keras/src/layers/preprocessing/benchmarks/category_vocab_file_dense_benchmark.py

@@ -1,110 +0,0 @@
-# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-"""Benchmark for KPL implementation of vocabulary columns from files with dense
-inputs."""
-
-import os
-
-import tensorflow.compat.v2 as tf
-
-import tf_keras.src as keras
-from tf_keras.src.layers.preprocessing import string_lookup
-from tf_keras.src.layers.preprocessing.benchmarks import (
-    feature_column_benchmark as fc_bm,
-)
-
-# isort: off
-from tensorflow.python.eager.def_function import (
-    function as tf_function,
-)
-
-NUM_REPEATS = 10
-BATCH_SIZES = [32, 256]
-
-
-class BenchmarkLayer(tf.test.TestCase, fc_bm.LayerBenchmark):
-    """Benchmark the layer forward pass."""
-
-    def _write_to_temp_file(self, file_name, vocab_list):
-        vocab_path = os.path.join(self.get_temp_dir(), file_name + ".txt")
-        with tf.io.gfile.GFile(vocab_path, "w") as writer:
-            for vocab in vocab_list:
-                writer.write(vocab + "\n")
-            writer.flush()
-            writer.close()
-        return vocab_path
-
-    def embedding_varlen(self, batch_size, max_length):
-        """Benchmark a variable-length embedding."""
-        # Data and constants.
-        vocab = fc_bm.create_vocabulary(32768)
-
-        path = self._write_to_temp_file("tmp", vocab)
-
-        data = fc_bm.create_string_data(
-            max_length, batch_size * NUM_REPEATS, vocab, pct_oov=0.15
-        )
-
-        # TF-Keras implementation
-        model = keras.Sequential()
-        model.add(
-            keras.Input(shape=(max_length,), name="data", dtype=tf.string)
-        )
-        model.add(string_lookup.StringLookup(vocabulary=path, mask_token=None))
-
-        # FC implementation
-        fc = tf.feature_column.categorical_column_with_vocabulary_list(
-            key="data", vocabulary_list=vocab, num_oov_buckets=1
-        )
-
-        # Wrap the FC implementation in a tf.function for a fair comparison
-        @tf_function()
-        def fc_fn(tensors):
-            fc.transform_feature(
-                tf.__internal__.feature_column.FeatureTransformationCache(
-                    tensors
-                ),
-                None,
-            )
-
-        # Benchmark runs
-        keras_data = {
-            "data": data.to_tensor(
-                default_value="", shape=(batch_size, max_length)
-            )
-        }
-        k_avg_time = fc_bm.run_keras(keras_data, model, batch_size, NUM_REPEATS)
-
-        fc_data = {
-            "data": data.to_tensor(
-                default_value="", shape=(batch_size, max_length)
-            )
-        }
-        fc_avg_time = fc_bm.run_fc(fc_data, fc_fn, batch_size, NUM_REPEATS)
-
-        return k_avg_time, fc_avg_time
-
-    def benchmark_layer(self):
-        for batch in BATCH_SIZES:
-            name = f"vocab_list|dense|batch_{batch}"
-            k_time, f_time = self.embedding_varlen(
-                batch_size=batch, max_length=256
-            )
-            self.report(name, k_time, f_time, NUM_REPEATS)
-
-
-if __name__ == "__main__":
-    tf.test.main()
-

tf_keras/src/layers/preprocessing/benchmarks/category_vocab_file_varlen_benchmark.py

@@ -1,103 +0,0 @@
-# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-"""Benchmark for KPL implementation of vocabulary columns from files with
-varying-length inputs."""
-
-import os
-
-import tensorflow.compat.v2 as tf
-
-import tf_keras.src as keras
-from tf_keras.src.layers.preprocessing import string_lookup
-from tf_keras.src.layers.preprocessing.benchmarks import (
-    feature_column_benchmark as fc_bm,
-)
-
-# isort: off
-from tensorflow.python.eager.def_function import (
-    function as tf_function,
-)
-
-NUM_REPEATS = 10
-BATCH_SIZES = [32, 256]
-
-
-class BenchmarkLayer(tf.test.TestCase, fc_bm.LayerBenchmark):
-    """Benchmark the layer forward pass."""
-
-    def _write_to_temp_file(self, file_name, vocab_list):
-        vocab_path = os.path.join(self.get_temp_dir(), file_name + ".txt")
-        with tf.io.gfile.GFile(vocab_path, "w") as writer:
-            for vocab in vocab_list:
-                writer.write(vocab + "\n")
-            writer.flush()
-            writer.close()
-        return vocab_path
-
-    def embedding_varlen(self, batch_size, max_length):
-        """Benchmark a variable-length embedding."""
-        # Data and constants.
-        vocab = fc_bm.create_vocabulary(32768)
-        path = self._write_to_temp_file("tmp", vocab)
-
-        data = fc_bm.create_string_data(
-            max_length, batch_size * NUM_REPEATS, vocab, pct_oov=0.15
-        )
-
-        # TF-Keras implementation
-        model = keras.Sequential()
-        model.add(
-            keras.Input(
-                shape=(max_length,), name="data", ragged=True, dtype=tf.string
-            )
-        )
-        model.add(string_lookup.StringLookup(vocabulary=path, mask_token=None))
-
-        # FC implementation
-        fc = tf.feature_column.sequence_categorical_column_with_vocabulary_list(
-            key="data", vocabulary_list=vocab, num_oov_buckets=1
-        )
-
-        # Wrap the FC implementation in a tf.function for a fair comparison
-        @tf_function()
-        def fc_fn(tensors):
-            fc.transform_feature(
-                tf.__internal__.feature_column.FeatureTransformationCache(
-                    tensors
-                ),
-                None,
-            )
-
-        # Benchmark runs
-        keras_data = {"data": data}
-        k_avg_time = fc_bm.run_keras(keras_data, model, batch_size, NUM_REPEATS)
-
-        fc_data = {"data": data.to_sparse()}
-        fc_avg_time = fc_bm.run_fc(fc_data, fc_fn, batch_size, NUM_REPEATS)
-
-        return k_avg_time, fc_avg_time
-
-    def benchmark_layer(self):
-        for batch in BATCH_SIZES:
-            name = f"vocab_list|varlen|batch_{batch}"
-            k_time, f_time = self.embedding_varlen(
-                batch_size=batch, max_length=256
-            )
-            self.report(name, k_time, f_time, NUM_REPEATS)
-
-
-if __name__ == "__main__":
-    tf.test.main()
-

tf_keras/src/layers/preprocessing/benchmarks/category_vocab_list_dense_benchmark.py

@@ -1,87 +0,0 @@
-# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-"""Benchmark for KPL implementation of vocabulary columns from lists with dense
-inputs."""
-
-import tensorflow.compat.v2 as tf
-
-import tf_keras.src as keras
-from tf_keras.src.layers.preprocessing import string_lookup
-from tf_keras.src.layers.preprocessing.benchmarks import (
-    feature_column_benchmark as fc_bm,
-)
-
-# isort: off
-from tensorflow.python.eager.def_function import (
-    function as tf_function,
-)
-
-NUM_REPEATS = 10
-BATCH_SIZES = [32, 256]
-
-
-def embedding_varlen(batch_size, max_length):
-    """Benchmark a variable-length embedding."""
-    # Data and constants.
-    vocab = fc_bm.create_vocabulary(32768)
-    data = fc_bm.create_string_data(
-        max_length, batch_size * NUM_REPEATS, vocab, pct_oov=0.15
-    )
-
-    # TF-Keras implementation
-    model = keras.Sequential()
-    model.add(keras.Input(shape=(max_length,), name="data", dtype=tf.string))
-    model.add(string_lookup.StringLookup(vocabulary=vocab, mask_token=None))
-
-    # FC implementation
-    fc = tf.feature_column.categorical_column_with_vocabulary_list(
-        key="data", vocabulary_list=vocab, num_oov_buckets=1
-    )
-
-    # Wrap the FC implementation in a tf.function for a fair comparison
-    @tf_function()
-    def fc_fn(tensors):
-        fc.transform_feature(
-            tf.__internal__.feature_column.FeatureTransformationCache(tensors),
-            None,
-        )
-
-    # Benchmark runs
-    keras_data = {
-        "data": data.to_tensor(default_value="", shape=(batch_size, max_length))
-    }
-    k_avg_time = fc_bm.run_keras(keras_data, model, batch_size, NUM_REPEATS)
-
-    fc_data = {
-        "data": data.to_tensor(default_value="", shape=(batch_size, max_length))
-    }
-    fc_avg_time = fc_bm.run_fc(fc_data, fc_fn, batch_size, NUM_REPEATS)
-
-    return k_avg_time, fc_avg_time
-
-
-class BenchmarkLayer(fc_bm.LayerBenchmark):
-    """Benchmark the layer forward pass."""
-
-    def benchmark_layer(self):
-        for batch in BATCH_SIZES:
-            name = f"vocab_list|dense|batch_{batch}"
-            k_time, f_time = embedding_varlen(batch_size=batch, max_length=256)
-            self.report(name, k_time, f_time, NUM_REPEATS)
-
-
-if __name__ == "__main__":
-    tf.test.main()
-

tf_keras/src/layers/preprocessing/benchmarks/category_vocab_list_indicator_dense_benchmark.py

@@ -1,96 +0,0 @@
-# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-"""Benchmark for KPL implementation of vocabulary columns + indicator from lists
-with dense inputs."""
-
-import tensorflow.compat.v2 as tf
-
-import tf_keras.src as keras
-from tf_keras.src.layers.preprocessing import category_encoding
-from tf_keras.src.layers.preprocessing import string_lookup
-from tf_keras.src.layers.preprocessing.benchmarks import (
-    feature_column_benchmark as fc_bm,
-)
-
-# isort: off
-from tensorflow.python.eager.def_function import (
-    function as tf_function,
-)
-
-NUM_REPEATS = 10
-BATCH_SIZES = [32, 256]
-
-
-def embedding_varlen(batch_size, max_length):
-    """Benchmark a variable-length embedding."""
-    # Data and constants.
-    vocab_size = 32768
-    vocab = fc_bm.create_vocabulary(vocab_size)
-    data = fc_bm.create_string_data(
-        max_length, batch_size * NUM_REPEATS, vocab, pct_oov=0.15
-    )
-
-    # TF-Keras implementation
-    model = keras.Sequential()
-    model.add(keras.Input(shape=(max_length,), name="data", dtype=tf.string))
-    model.add(string_lookup.StringLookup(vocabulary=vocab, mask_token=None))
-    model.add(
-        category_encoding.CategoryEncoding(
-            num_tokens=vocab_size + 1, output_mode="count"
-        )
-    )
-
-    # FC implementation
-    fc = tf.feature_column.indicator_column(
-        tf.feature_column.categorical_column_with_vocabulary_list(
-            key="data", vocabulary_list=vocab, num_oov_buckets=1
-        )
-    )
-
-    # Wrap the FC implementation in a tf.function for a fair comparison
-    @tf_function()
-    def fc_fn(tensors):
-        fc.transform_feature(
-            tf.__internal__.feature_column.FeatureTransformationCache(tensors),
-            None,
-        )
-
-    # Benchmark runs
-    keras_data = {
-        "data": data.to_tensor(default_value="", shape=(batch_size, max_length))
-    }
-    k_avg_time = fc_bm.run_keras(keras_data, model, batch_size, NUM_REPEATS)
-
-    fc_data = {
-        "data": data.to_tensor(default_value="", shape=(batch_size, max_length))
-    }
-    fc_avg_time = fc_bm.run_fc(fc_data, fc_fn, batch_size, NUM_REPEATS)
-
-    return k_avg_time, fc_avg_time
-
-
-class BenchmarkLayer(fc_bm.LayerBenchmark):
-    """Benchmark the layer forward pass."""
-
-    def benchmark_layer(self):
-        for batch in BATCH_SIZES:
-            name = f"vocab_list_indicator|dense|batch_{batch}"
-            k_time, f_time = embedding_varlen(batch_size=batch, max_length=256)
-            self.report(name, k_time, f_time, NUM_REPEATS)
-
-
-if __name__ == "__main__":
-    tf.test.main()
-

tf_keras/src/layers/preprocessing/benchmarks/category_vocab_list_indicator_varlen_benchmark.py

@@ -1,96 +0,0 @@
-# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-"""Benchmark for KPL implementation of vocabulary columns + indicator from lists
-with varying-length inputs."""
-
-import tensorflow.compat.v2 as tf
-
-import tf_keras.src as keras
-from tf_keras.src.layers.preprocessing import category_encoding
-from tf_keras.src.layers.preprocessing import string_lookup
-from tf_keras.src.layers.preprocessing.benchmarks import (
-    feature_column_benchmark as fc_bm,
-)
-
-# isort: off
-from tensorflow.python.eager.def_function import (
-    function as tf_function,
-)
-
-NUM_REPEATS = 10
-BATCH_SIZES = [32, 256]
-
-
-def embedding_varlen(batch_size, max_length):
-    """Benchmark a variable-length embedding."""
-    # Data and constants.
-    vocab_size = 32768
-    vocab = fc_bm.create_vocabulary(vocab_size)
-    data = fc_bm.create_string_data(
-        max_length, batch_size * NUM_REPEATS, vocab, pct_oov=0.15
-    )
-
-    # TF-Keras implementation
-    model = keras.Sequential()
-    model.add(
-        keras.Input(
-            shape=(max_length,), name="data", ragged=True, dtype=tf.string
-        )
-    )
-    model.add(string_lookup.StringLookup(vocabulary=vocab, mask_token=None))
-    model.add(
-        category_encoding.CategoryEncoding(
-            num_tokens=vocab_size + 1, output_mode="count"
-        )
-    )
-
-    # FC implementation
-    fc = tf.feature_column.indicator_column(
-        tf.feature_column.sequence_categorical_column_with_vocabulary_list(
-            key="data", vocabulary_list=vocab, num_oov_buckets=1
-        )
-    )
-
-    # Wrap the FC implementation in a tf.function for a fair comparison
-    @tf_function()
-    def fc_fn(tensors):
-        fc.transform_feature(
-            tf.__internal__.feature_column.FeatureTransformationCache(tensors),
-            None,
-        )
-
-    # Benchmark runs
-    keras_data = {"data": data}
-    k_avg_time = fc_bm.run_keras(keras_data, model, batch_size, NUM_REPEATS)
-
-    fc_data = {"data": data.to_sparse()}
-    fc_avg_time = fc_bm.run_fc(fc_data, fc_fn, batch_size, NUM_REPEATS)
-
-    return k_avg_time, fc_avg_time
-
-
-class BenchmarkLayer(fc_bm.LayerBenchmark):
-    """Benchmark the layer forward pass."""
-
-    def benchmark_layer(self):
-        for batch in BATCH_SIZES:
-            name = f"vocab_list_indicator|varlen|batch_{batch}"
-            k_time, f_time = embedding_varlen(batch_size=batch, max_length=256)
-            self.report(name, k_time, f_time, NUM_REPEATS)
-
-
-if __name__ == "__main__":
-    tf.test.main()
-

tf_keras/src/layers/preprocessing/benchmarks/category_vocab_list_varlen_benchmark.py

@@ -1,87 +0,0 @@
-# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-"""Benchmark for KPL implementation of vocabulary columns from lists with
-varying-length inputs."""
-
-import tensorflow.compat.v2 as tf
-
-import tf_keras.src as keras
-from tf_keras.src.layers.preprocessing import string_lookup
-from tf_keras.src.layers.preprocessing.benchmarks import (
-    feature_column_benchmark as fc_bm,
-)
-
-# isort: off
-from tensorflow.python.eager.def_function import (
-    function as tf_function,
-)
-
-NUM_REPEATS = 10
-BATCH_SIZES = [32, 256]
-
-
-def embedding_varlen(batch_size, max_length):
-    """Benchmark a variable-length embedding."""
-    # Data and constants.
-    vocab = fc_bm.create_vocabulary(32768)
-    data = fc_bm.create_string_data(
-        max_length, batch_size * NUM_REPEATS, vocab, pct_oov=0.15
-    )
-
-    # TF-Keras implementation
-    model = keras.Sequential()
-    model.add(
-        keras.Input(
-            shape=(max_length,), name="data", ragged=True, dtype=tf.string
-        )
-    )
-    model.add(string_lookup.StringLookup(vocabulary=vocab, mask_token=None))
-
-    # FC implementation
-    fc = tf.feature_column.sequence_categorical_column_with_vocabulary_list(
-        key="data", vocabulary_list=vocab, num_oov_buckets=1
-    )
-
-    # Wrap the FC implementation in a tf.function for a fair comparison
-    @tf_function()
-    def fc_fn(tensors):
-        fc.transform_feature(
-            tf.__internal__.feature_column.FeatureTransformationCache(tensors),
-            None,
-        )
-
-    # Benchmark runs
-    keras_data = {"data": data}
-    k_avg_time = fc_bm.run_keras(keras_data, model, batch_size, NUM_REPEATS)
-
-    fc_data = {"data": data.to_sparse()}
-    fc_avg_time = fc_bm.run_fc(fc_data, fc_fn, batch_size, NUM_REPEATS)
-
-    return k_avg_time, fc_avg_time
-
-
-class BenchmarkLayer(fc_bm.LayerBenchmark):
-    """Benchmark the layer forward pass."""
-
-    def benchmark_layer(self):
-        for batch in BATCH_SIZES:
-            name = f"vocab_list|varlen|batch_{batch}"
-            k_time, f_time = embedding_varlen(batch_size=batch, max_length=256)
-            self.report(name, k_time, f_time, NUM_REPEATS)
-
-
-if __name__ == "__main__":
-    tf.test.main()
-

tf_keras/src/layers/preprocessing/benchmarks/discretization_adapt_benchmark.py

@@ -1,109 +0,0 @@
-# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-"""Benchmark for TF-Keras discretization preprocessing layer's adapt method."""
-
-import time
-
-import numpy as np
-import tensorflow.compat.v2 as tf
-
-import tf_keras.src as keras
-from tf_keras.src.layers.preprocessing import discretization
-
-EPSILON = 0.1
-
-
-def reduce_fn(state, values, epsilon=EPSILON):
-    """tf.data.Dataset-friendly implementation of mean and variance."""
-
-    (state_,) = state
-    summary = discretization.summarize(values, epsilon)
-    if np.sum(state_[:, 0]) == 0:
-        return (summary,)
-    return (discretization.merge_summaries(state_, summary, epsilon),)
-
-
-class BenchmarkAdapt(tf.test.Benchmark):
-    """Benchmark adapt."""
-
-    def run_dataset_implementation(self, num_elements, batch_size):
-        input_t = keras.Input(shape=(1,))
-        layer = discretization.Discretization()
-        _ = layer(input_t)
-
-        num_repeats = 5
-        starts = []
-        ends = []
-        for _ in range(num_repeats):
-            ds = tf.data.Dataset.range(num_elements)
-            ds = ds.map(lambda x: tf.expand_dims(tf.cast(x, tf.float32), -1))
-            ds = ds.batch(batch_size)
-
-            starts.append(time.time())
-            # Benchmarked code begins here.
-            state = ds.reduce((np.zeros((1, 2)),), reduce_fn)
-
-            bins = discretization.get_bucket_boundaries(state, 100)
-            layer.set_weights([bins])
-            # Benchmarked code ends here.
-            ends.append(time.time())
-
-        avg_time = np.mean(np.array(ends) - np.array(starts))
-        return avg_time
-
-    def bm_adapt_implementation(self, num_elements, batch_size):
-        """Test the KPL adapt implementation."""
-        input_t = keras.Input(shape=(1,), dtype=tf.float32)
-        layer = discretization.Discretization()
-        _ = layer(input_t)
-
-        num_repeats = 5
-        starts = []
-        ends = []
-        for _ in range(num_repeats):
-            ds = tf.data.Dataset.range(num_elements)
-            ds = ds.map(lambda x: tf.expand_dims(tf.cast(x, tf.float32), -1))
-            ds = ds.batch(batch_size)
-
-            starts.append(time.time())
-            # Benchmarked code begins here.
-            layer.adapt(ds)
-            # Benchmarked code ends here.
-            ends.append(time.time())
-
-        avg_time = np.mean(np.array(ends) - np.array(starts))
-        name = "discretization_adapt|%s_elements|batch_%s" % (
-            num_elements,
-            batch_size,
-        )
-        baseline = self.run_dataset_implementation(num_elements, batch_size)
-        extras = {
-            "tf.data implementation baseline": baseline,
-            "delta seconds": (baseline - avg_time),
-            "delta percent": ((baseline - avg_time) / baseline) * 100,
-        }
-        self.report_benchmark(
-            iters=num_repeats, wall_time=avg_time, extras=extras, name=name
-        )
-
-    def benchmark_vocab_size_by_batch(self):
-        for vocab_size in [100, 1000, 10000, 100000, 1000000]:
-            for batch in [64 * 2048]:
-                self.bm_adapt_implementation(vocab_size, batch)
-
-
-if __name__ == "__main__":
-    tf.test.main()
-