Update keras benchmark test to use public TF API.

PiperOrigin-RevId: 339489416
Change-Id: If526ae1f84e3516c85c683a823ad85d7e36baeb8

tensorflow/python/keras/layers/preprocessing/benchmarks/category_crossing_benchmark.py

@@ -45,7 +45,7 @@ def int_gen():
     yield (np.random.randint(0, 5, (1,)), np.random.randint(0, 7, (1,)))
 
 
-class BenchmarkLayer(benchmark.Benchmark):
+class BenchmarkLayer(benchmark.TensorFlowBenchmark):
   """Benchmark the layer forward pass."""
 
   def run_dataset_implementation(self, batch_size):

tensorflow/python/keras/layers/preprocessing/benchmarks/category_encoding_benchmark.py

@@ -36,7 +36,7 @@ FLAGS = flags.FLAGS
 v2_compat.enable_v2_behavior()
 
 
-class BenchmarkLayer(benchmark.Benchmark):
+class BenchmarkLayer(benchmark.TensorFlowBenchmark):
   """Benchmark the layer forward pass."""
 
   def run_dataset_implementation(self, output_mode, batch_size, sequence_length,

tensorflow/python/keras/layers/preprocessing/benchmarks/discretization_adapt_benchmark.py

@@ -48,7 +48,7 @@ def reduce_fn(state, values, epsilon=EPSILON):
   return (discretization.merge_summaries(state_, summary, epsilon),)
 
 
-class BenchmarkAdapt(benchmark.Benchmark):
+class BenchmarkAdapt(benchmark.TensorFlowBenchmark):
   """Benchmark adapt."""
 
   def run_dataset_implementation(self, num_elements, batch_size):

tensorflow/python/keras/layers/preprocessing/benchmarks/hashing_benchmark.py

@@ -47,7 +47,7 @@ def word_gen():
     yield "".join(random.choice(string.ascii_letters) for i in range(2))
 
 
-class BenchmarkLayer(benchmark.Benchmark):
+class BenchmarkLayer(benchmark.TensorFlowBenchmark):
   """Benchmark the layer forward pass."""
 
   def run_dataset_implementation(self, batch_size):

tensorflow/python/keras/layers/preprocessing/benchmarks/image_preproc_benchmark.py

@@ -85,7 +85,7 @@ def image_augmentation(inputs, batch_size):
   return img
 
 
-class BenchmarkLayer(benchmark.Benchmark):
+class BenchmarkLayer(benchmark.TensorFlowBenchmark):
   """Benchmark the layer forward pass."""
 
   def run_dataset_implementation(self, batch_size):

tensorflow/python/keras/layers/preprocessing/benchmarks/index_lookup_adapt_benchmark.py

@@ -63,7 +63,7 @@ def get_top_k(dataset, k):
   return sorted_vocab
 
 
-class BenchmarkAdapt(benchmark.Benchmark):
+class BenchmarkAdapt(benchmark.TensorFlowBenchmark):
   """Benchmark adapt."""
 
   def run_numpy_implementation(self, num_elements, batch_size, k):

tensorflow/python/keras/layers/preprocessing/benchmarks/normalization_adapt_benchmark.py

@@ -60,7 +60,7 @@ def reduce_fn(state, values):
   return (k, n + batch_size, ex, ex2)
 
 
-class BenchmarkAdapt(benchmark.Benchmark):
+class BenchmarkAdapt(benchmark.TensorFlowBenchmark):
   """Benchmark adapt."""
 
   def run_dataset_implementation(self, num_elements, batch_size):