Allow benchmarks and tests namespaces in pip package.

PiperOrigin-RevId: 512157283

keras/benchmarks/benchmark_util.py

@@ -17,8 +17,8 @@
 import timeit
 
 import numpy as np
-import tensorflow.compat.v2 as tf
 
+from keras import callbacks
 from keras.benchmarks import distribution_util
 
 
@@ -72,7 +72,7 @@ def get_keras_examples_metadata(
     }
 
 
-class TimerCallBack(tf.keras.callbacks.Callback):
+class TimerCallBack(callbacks.Callback):
     """Callback for logging time in each epoch or batch."""
 
     def __init__(self):

keras/benchmarks/model_memory_profile.py

@@ -21,11 +21,12 @@
 """
 
 import numpy as np
-import tensorflow.compat.v2 as tf
 from absl import app
 from absl import flags
 from absl import logging
 
+import keras
+
 try:
     import memory_profiler
 except ImportError:
@@ -43,10 +44,10 @@ def _imdb_lstm_model():
     y_train = np.random.random((2500, 1))
 
     # IMDB LSTM model.
-    model = tf.keras.Sequential()
-    model.add(tf.keras.layers.Embedding(20000, 128))
-    model.add(tf.keras.layers.LSTM(128, dropout=0.2, recurrent_dropout=0.2))
-    model.add(tf.keras.layers.Dense(1, activation="sigmoid"))
+    model = keras.Sequential()
+    model.add(keras.layers.Embedding(20000, 128))
+    model.add(keras.layers.LSTM(128, dropout=0.2, recurrent_dropout=0.2))
+    model.add(keras.layers.Dense(1, activation="sigmoid"))
 
     model.compile("sgd", "mse")
     # Warm up the model with one epoch.

keras/benchmarks/saved_model_benchmarks/saved_model_benchmark_util.py

@@ -23,6 +23,8 @@
 
 import tensorflow.compat.v2 as tf
 
+import keras
+
 
 def save_and_load_benchmark(app):
     """Util for saved model benchmarks."""
@@ -40,15 +42,15 @@ def save_and_load_benchmark(app):
 
     # Run one untimed iteration of saving/loading.
     model.save(save_dir, save_format="tf")
-    tf.keras.models.load_model(save_dir)
+    keras.models.load_model(save_dir)
 
     for _ in range(trials):
         start_time = time.time()
         model.save(save_dir, save_format="tf")
         total_save_time += time.time() - start_time
 
         start_time = time.time()
-        tf.keras.models.load_model(save_dir)
+        keras.models.load_model(save_dir)
         total_load_time += time.time() - start_time
 
     save_result = {

pip_build.py

@@ -54,14 +54,12 @@
 VERBOSE = True
 INIT_FILE_HEADER = """AUTOGENERATED. DO NOT EDIT."""
 # These are symbols that have export issues and that we skip for now.
-SYMBOLS_TO_SKIP = ["layer_test"]
+SYMBOLS_TO_SKIP = []
 
 
 def copy_keras_codebase(source_dir, target_dir):
     disallowed = [
-        "benchmarks",
         "tools",
-        "tests",
         "integration_test",
     ]