correct pylint formattings - 2

tensorflow · Jun 4, 2021 · 85ca9b2 · 85ca9b2
1 parent 7ba767d
commit 85ca9b2
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 15 deletions.
diff --git a/tensorflow/python/distribute/input_lib.py b/tensorflow/python/distribute/input_lib.py
@@ -2002,10 +2002,11 @@ def _get_dataset_attributes(dataset):
 
 def _should_use_multi_device_iterator(options):
   """Determine whether to use multi_device_iterator_ops.OwnedMultiDeviceIterator"""
-  if (options is None
-      or options.experimental_replication_mode == InputReplicationMode.PER_WORKER
-      or (options.experimental_replication_mode == InputReplicationMode.PER_REPLICA
-          and options.experimental_prefetch_to_device)):
+  if (options is None or
+      options.experimental_replication_mode == InputReplicationMode.PER_WORKER
+      or
+      (options.experimental_replication_mode == InputReplicationMode.PER_REPLICA
+        and options.experimental_prefetch_to_device)):
     return True
   return False
 

diff --git a/tensorflow/python/distribute/input_lib_type_spec_test.py b/tensorflow/python/distribute/input_lib_type_spec_test.py
@@ -462,15 +462,16 @@ def f(v):
           enable_get_next_as_optional=[True, False],
           experimental_place_dataset_on_device=[True, False],
           experimental_prefetch_to_device=[True, False],))
-  def testFromFunctionInputSignatureForPerReplicaValuesWithOptions(self, distribution,
-                                                                   enable_get_next_as_optional,
-                                                                   experimental_place_dataset_on_device,
-                                                                   experimental_prefetch_to_device):
-    
+  def testFromFunctionInputSignatureForPerReplicaValuesWithOptions(
+      self, distribution, enable_get_next_as_optional,
+      experimental_place_dataset_on_device,
+      experimental_prefetch_to_device):
+
     if experimental_place_dataset_on_device and experimental_prefetch_to_device:
       self.skipTest("Setting experimental_place_dataset_on_device and "
-                    "experimental_prefetch_to_device to `True` is not allowed "
-                    "when using distribute_lib.InputReplicationMode.PER_REPLICA.")
+                    "experimental_prefetch_to_device to `True` is not "
+                    "allowed when using "
+                    "distribute_lib.InputReplicationMode.PER_REPLICA.")
 
     fname1 = os.path.join(self.get_temp_dir(), "1.txt")
     _create_text_file(fname1, 5)
@@ -486,13 +487,16 @@ def dataset_fn(input_context):
               input_context.get_per_replica_batch_size(4))
 
     options = distribute_lib.InputOptions(
-        experimental_place_dataset_on_device=experimental_place_dataset_on_device,
-        experimental_prefetch_to_device=experimental_prefetch_to_device,
-        experimental_replication_mode=distribute_lib.InputReplicationMode.PER_REPLICA) 
+      experimental_place_dataset_on_device=experimental_place_dataset_on_device,
+      experimental_prefetch_to_device=experimental_prefetch_to_device,
+      experimental_replication_mode=(
+        distribute_lib.InputReplicationMode.PER_REPLICA
+      ))
 
     distribution.extended.experimental_enable_get_next_as_optional = (
         enable_get_next_as_optional)
-    ds = distribution.experimental_distribute_datasets_from_function(dataset_fn, options)
+    ds = distribution.experimental_distribute_datasets_from_function(dataset_fn,
+                                                                     options)
 
     iterator = iter(ds)
     _check_type_spec_structure(iterator)