Formats with shell/format.sh

PiperOrigin-RevId: 564465973

pip_build.py

@@ -77,7 +77,10 @@ def ignore(path, names):
 
 def convert_keras_imports(src_directory):
     def _convert_line(line):
-        if "import tf_keras.protobuf" in line or "from tf_keras.protobuf" in line:
+        if (
+            "import tf_keras.protobuf" in line
+            or "from tf_keras.protobuf" in line
+        ):
             return line
         # Imports starting from `root_name`.
         if line.strip() == f"import {PACKAGE_NAME}":
@@ -100,7 +103,9 @@ def _convert_line(line):
             f"from {PACKAGE_NAME}.{SRC_DIRNAME} import",
         )
         # A way to catch LazyLoader calls. Hacky.
-        line = line.replace('globals(), "tf_keras.', 'globals(), "tf_keras.src.')
+        line = line.replace(
+            'globals(), "tf_keras.', 'globals(), "tf_keras.src.'
+        )
         return line
 
     for root, _, files in os.walk(src_directory):

tf_keras/api/extractor_wrapper.py

@@ -14,7 +14,9 @@
 # ==============================================================================
 """Thin wrapper to call TensorFlow's API extractor script."""
 from absl import app
-from tensorflow.python.tools.api.generator2.extractor import extractor
+from tensorflow.python.tools.api.generator2.extractor import (
+    extractor,
+)
 
 if __name__ == "__main__":
     app.run(extractor.main)

tf_keras/api/generator_wrapper.py

@@ -14,7 +14,9 @@
 # ==============================================================================
 """Thin wrapper to call TensorFlow's API generator script."""
 from absl import app
-from tensorflow.python.tools.api.generator2.generator import generator
+from tensorflow.python.tools.api.generator2.generator import (
+    generator,
+)
 
 if __name__ == "__main__":
     app.run(generator.main)

tf_keras/applications/mobilenet_v2.py

@@ -254,7 +254,6 @@ def MobileNetV2(
 
     # If input_shape is None, infer shape from input_tensor.
     if input_shape is None and input_tensor is not None:
-
         try:
             backend.is_keras_tensor(input_tensor)
         except ValueError:

tf_keras/applications/mobilenet_v3.py

@@ -247,7 +247,6 @@ def MobileNetV3(
 
     # If input_shape is None, infer shape from input_tensor
     if input_shape is None and input_tensor is not None:
-
         try:
             backend.is_keras_tensor(input_tensor)
         except ValueError:

tf_keras/benchmarks/keras_examples_benchmarks/mnist_conv_custom_training_benchmark_test.py

@@ -172,7 +172,6 @@ def train_function(
         for _ in range(epochs):
             # Iterate over the batches of the dataset.
             for batch_dataset in train_dataset:
-
                 start_time = timer()
 
                 if distribution_strategy is not None:

tf_keras/benchmarks/layer_benchmarks/layer_benchmarks_test.py

@@ -431,7 +431,6 @@ class KerasLayerBenchmarks(
     layer_benchmarks_test_base.LayerBenchmarksBase,
     metaclass=tf.__internal__.test.ParameterizedBenchmark,
 ):
-
     # The parameter of each layer benchmark is a tuple, and the first one is
     # the benchmark name. It must follow the convention of
     # "{layer_name}_{small|normal|large}_shape" to make it compatible with

tf_keras/benchmarks/metrics_memory_benchmark_test.py

@@ -24,7 +24,6 @@
 
 
 class KerasMetricMemoryBenchmark(tf.test.Benchmark):
-
     # This test is added to measure the memory footprint for
     # metrics_utils._update_confusion_matrix_variables_optimized().
 

tf_keras/benchmarks/saved_model_benchmarks/densenet_benchmark_test.py

@@ -20,7 +20,9 @@
 
 import tensorflow.compat.v2 as tf
 
-from tf_keras.benchmarks.saved_model_benchmarks import saved_model_benchmark_util
+from tf_keras.benchmarks.saved_model_benchmarks import (
+    saved_model_benchmark_util,
+)
 
 
 class BenchmarkSaveApplications(tf.test.Benchmark):

tf_keras/benchmarks/saved_model_benchmarks/efficientnet_benchmark_test.py

@@ -20,7 +20,9 @@
 
 import tensorflow.compat.v2 as tf
 
-from tf_keras.benchmarks.saved_model_benchmarks import saved_model_benchmark_util
+from tf_keras.benchmarks.saved_model_benchmarks import (
+    saved_model_benchmark_util,
+)
 
 
 class BenchmarkSaveApplications(tf.test.Benchmark):

tf_keras/benchmarks/saved_model_benchmarks/inception_resnet_v2_benchmark_test.py

@@ -20,7 +20,9 @@
 
 import tensorflow.compat.v2 as tf
 
-from tf_keras.benchmarks.saved_model_benchmarks import saved_model_benchmark_util
+from tf_keras.benchmarks.saved_model_benchmarks import (
+    saved_model_benchmark_util,
+)
 
 
 class BenchmarkSaveApplications(tf.test.Benchmark):

tf_keras/benchmarks/saved_model_benchmarks/mobilenet_benchmark_test.py

@@ -20,7 +20,9 @@
 
 import tensorflow.compat.v2 as tf
 
-from tf_keras.benchmarks.saved_model_benchmarks import saved_model_benchmark_util
+from tf_keras.benchmarks.saved_model_benchmarks import (
+    saved_model_benchmark_util,
+)
 
 
 class BenchmarkSaveApplications(tf.test.Benchmark):

tf_keras/benchmarks/saved_model_benchmarks/nasnet_large_benchmark_test.py

@@ -20,7 +20,9 @@
 
 import tensorflow.compat.v2 as tf
 
-from tf_keras.benchmarks.saved_model_benchmarks import saved_model_benchmark_util
+from tf_keras.benchmarks.saved_model_benchmarks import (
+    saved_model_benchmark_util,
+)
 
 
 class BenchmarkSaveApplications(tf.test.Benchmark):

tf_keras/benchmarks/saved_model_benchmarks/resnet152_v2_benchmark_test.py

@@ -20,7 +20,9 @@
 
 import tensorflow.compat.v2 as tf
 
-from tf_keras.benchmarks.saved_model_benchmarks import saved_model_benchmark_util
+from tf_keras.benchmarks.saved_model_benchmarks import (
+    saved_model_benchmark_util,
+)
 
 
 class BenchmarkSaveApplications(tf.test.Benchmark):

tf_keras/benchmarks/saved_model_benchmarks/vgg_benchmark_test.py

@@ -20,7 +20,9 @@
 
 import tensorflow.compat.v2 as tf
 
-from tf_keras.benchmarks.saved_model_benchmarks import saved_model_benchmark_util
+from tf_keras.benchmarks.saved_model_benchmarks import (
+    saved_model_benchmark_util,
+)
 
 
 class BenchmarkSaveApplications(tf.test.Benchmark):

tf_keras/benchmarks/saved_model_benchmarks/xception_benchmark_test.py

@@ -20,7 +20,9 @@
 
 import tensorflow.compat.v2 as tf
 
-from tf_keras.benchmarks.saved_model_benchmarks import saved_model_benchmark_util
+from tf_keras.benchmarks.saved_model_benchmarks import (
+    saved_model_benchmark_util,
+)
 
 
 class BenchmarkSaveApplications(tf.test.Benchmark):

tf_keras/callbacks_v1.py

@@ -420,7 +420,6 @@ def on_epoch_begin(self, epoch, logs=None):
 
         # check if histogram summary should be run for this epoch
         if self.histogram_freq and epoch % self.histogram_freq == 0:
-
             # add the histogram summary op if it should run this epoch
             self.model._make_test_function()
             if self.merged not in self.model.test_function.fetches:
@@ -448,7 +447,6 @@ def on_epoch_end(self, epoch, logs=None):
 
         # pop the histogram summary op after each epoch
         if self.histogram_freq:
-
             if self.merged in self.model.test_function.fetches:
                 self.model.test_function.fetches.remove(self.merged)
             if self.merged in self.model.test_function.fetch_callbacks:

tf_keras/distribute/collective_all_reduce_strategy_test.py

@@ -19,7 +19,9 @@
 
 from tf_keras import layers
 from tf_keras.engine import training
-from tf_keras.optimizers.legacy import gradient_descent as gradient_descent_keras
+from tf_keras.optimizers.legacy import (
+    gradient_descent as gradient_descent_keras,
+)
 from tf_keras.testing_infra import test_utils
 
 

tf_keras/distribute/custom_training_loop_models_test.py

@@ -205,7 +205,6 @@ def step_fn(inputs):
         train_step(input_iterator)
 
     def test_lstm(self, distribution):
-
         batch_size = 32
 
         def create_lstm_model():
@@ -481,7 +480,6 @@ def test_variable_run_argument(self, distribution):
         dataset = dataset.batch(2, drop_remainder=True)
 
         def replica_step(trainable_variables, features):
-
             with tf.GradientTape() as tape:
                 net_out = net(features[0], training=True)
                 loss = (net_out - 1.0) * (net_out - 1.0)

tf_keras/distribute/custom_training_loop_optimizer_test.py

@@ -46,7 +46,6 @@ class OptimizerTest(tf.test.TestCase, parameterized.TestCase):
     def test_custom_aggregation(
         self, distribution, experimental_aggregate_gradients, expected
     ):
-
         with distribution.scope():
             v = tf.Variable([0.0, 0.0])
             optimizer = gradient_descent.SGD(0.1)
@@ -90,7 +89,6 @@ def step_fn(grads):
     def test_custom_aggregation_one_device(
         self, distribution, experimental_aggregate_gradients
     ):
-
         with distribution.scope():
             v = tf.Variable([0.0, 0.0])
             optimizer = gradient_descent.SGD(0.1)

tf_keras/distribute/dataset_creator_model_fit_test_base.py

@@ -36,9 +36,7 @@ class DatasetCreatorModelFitTestBase(tf.test.TestCase, parameterized.TestCase):
     """The base class for DatasetCreator with Model.fit tests."""
 
     def _get_dataset_fn(self, use_lookup_layer):
-
         if use_lookup_layer:
-
             filepath = os.path.join(self.get_temp_dir(), "vocab")
             with open(filepath, "w") as f:
                 f.write("\n".join(["earth", "wind", "and", "fire"]))

tf_keras/distribute/distribute_coordinator_utils.py

@@ -137,7 +137,6 @@ def __enter__(self):
     def __exit__(
         self, unused_exception_type, unused_exception_value, unused_traceback
     ):
-
         _worker_context.current = None
 
     def _get_master_target(self):
@@ -465,7 +464,6 @@ def join(self):
 def _configure_session_config_for_std_servers(
     strategy, eval_strategy, session_config, cluster_spec, task_type, task_id
 ):
-
     """Call strategy's `configure` to mutate the session_config.
 
     The session_config is currently needed as default config for a TensorFlow

tf_keras/distribute/distribute_strategy_test.py

@@ -38,7 +38,9 @@
 from tf_keras.engine import base_layer_utils
 from tf_keras.mixed_precision import policy
 from tf_keras.optimizers import optimizer as optimizer_base
-from tf_keras.optimizers.legacy import gradient_descent as gradient_descent_keras
+from tf_keras.optimizers.legacy import (
+    gradient_descent as gradient_descent_keras,
+)
 from tf_keras.testing_infra import test_utils
 from tf_keras.utils import losses_utils
 from tf_keras.utils import np_utils
@@ -657,7 +659,6 @@ def run_fn():
     def test_optimizer_in_cross_replica_context_raises_error(
         self, distribution
     ):
-
         with self.cached_session(), distribution.scope():
             model = keras.models.Sequential([keras.layers.Dense(1)])
             x = np.array([[1.0]])
@@ -949,7 +950,6 @@ def test_predict_multi_output_model_with_partial_batch(self, distribution):
         all_strategy_combinations()
     )
     def test_gradients_are_none(self, distribution):
-
         if not tf.executing_eagerly():
             self.skipTest("None gradients are not supported in graph mode")
 
@@ -1127,7 +1127,6 @@ def test_fit_with_tuple_and_dict_dataset_inputs(self, distribution):
         all_strategy_combinations()
     )
     def test_fit_with_dictionary_in_the_dataset_b135161171(self, distribution):
-
         if backend.is_tpu_strategy(distribution):
             self.skipTest("b/142805125")
 
@@ -1228,7 +1227,6 @@ def test_fit_eval_and_predict_methods_on_dataset_without_steps(
     def test_predict_on_dataset_with_unknown_cardinality_without_steps(
         self, distribution, mode
     ):
-
         if mode == "graph" and backend.is_tpu_strategy(distribution):
             self.skipTest("partial batch not supported with TPU in graph mode.")
 
@@ -1419,9 +1417,7 @@ def test_fit_eval_and_predict_methods_on_dataset(self, distribution):
     )
     def test_fit_eval_and_predict_with_optimizer(self, distribution, optimizer):
         with self.cached_session():
-
             with distribution.scope():
-
                 model = get_model()
                 loss = "mse"
                 model.compile(optimizer(), loss)

tf_keras/distribute/distributed_training_utils.py

@@ -131,7 +131,6 @@ def get_strategy():
 
 
 def maybe_preemption_handler_scope(model):
-
     if getattr(model, "_preemption_handler", None):
         preemption_checkpoint_scope = (
             model._preemption_handler.watch_preemption_scope()

tf_keras/distribute/distributed_training_utils_v1.py

@@ -264,7 +264,6 @@ def validate_callbacks(input_callbacks, optimizer):
                 callback,
                 (callbacks.LearningRateScheduler, callbacks.ReduceLROnPlateau),
             ):
-
                 if not isinstance(optimizer, optimizer_v2.OptimizerV2):
                     raise ValueError(
                         "You must specify a TF-Keras Optimizer V2 when using "

tf_keras/distribute/keras_dnn_correctness_test.py

@@ -21,7 +21,9 @@
 from tf_keras import backend
 from tf_keras.distribute import keras_correctness_test_base
 from tf_keras.distribute import strategy_combinations
-from tf_keras.optimizers.legacy import gradient_descent as gradient_descent_keras
+from tf_keras.optimizers.legacy import (
+    gradient_descent as gradient_descent_keras,
+)
 from tf_keras.testing_infra import test_utils
 
 

tf_keras/distribute/keras_embedding_model_correctness_test.py

@@ -19,7 +19,9 @@
 
 import tf_keras as keras
 from tf_keras.distribute import keras_correctness_test_base
-from tf_keras.optimizers.legacy import gradient_descent as gradient_descent_keras
+from tf_keras.optimizers.legacy import (
+    gradient_descent as gradient_descent_keras,
+)
 
 
 class DistributionStrategyEmbeddingModelCorrectnessTest(
@@ -65,7 +67,6 @@ def get_model(
     def test_embedding_model_correctness(
         self, distribution, use_numpy, use_validation_data
     ):
-
         self.use_distributed_dense = False
         self.run_correctness_test(distribution, use_numpy, use_validation_data)
 

tf_keras/distribute/keras_rnn_model_correctness_test.py

@@ -24,7 +24,9 @@
 from tf_keras.layers.rnn import lstm
 from tf_keras.layers.rnn import lstm_v1
 from tf_keras.mixed_precision import policy
-from tf_keras.optimizers.legacy import gradient_descent as gradient_descent_keras
+from tf_keras.optimizers.legacy import (
+    gradient_descent as gradient_descent_keras,
+)
 from tf_keras.testing_infra import test_utils
 
 

tf_keras/distribute/keras_stateful_lstm_model_correctness_test.py

@@ -19,7 +19,9 @@
 
 import tf_keras as keras
 from tf_keras.distribute import keras_correctness_test_base
-from tf_keras.optimizers.legacy import gradient_descent as gradient_descent_keras
+from tf_keras.optimizers.legacy import (
+    gradient_descent as gradient_descent_keras,
+)
 
 
 def strategies_for_stateful_embedding_model():

tf_keras/distribute/keras_utils_test.py

@@ -400,7 +400,6 @@ def test_standalone_loss_without_loss_reduction(self, distribution):
 class TestDistributionStrategyWithLossMasking(
     tf.test.TestCase, parameterized.TestCase
 ):
-
     # TODO(priyag): Enable all strategies for this test. Currently it does not
     # work for TPU due to some invalid datatype.
     @tf.__internal__.distribute.combinations.generate(

tf_keras/distribute/mirrored_variable_test.py

@@ -86,7 +86,6 @@ def _is_mirrored(self, val):
 
     def testWithLayers(self, distribution):
         def model_fn(features):
-
             layer1 = core.Dense(1)
             layer1(features)
             layer2 = core.Dense(1)