apply Black 2024 style in fbcode (4/16)

Summary:
Formats the covered files with pyfmt.

paintitblack

Reviewed By: aleivag

Differential Revision: D54447727

fbshipit-source-id: 8844b1caa08de94d04ac4df3c768dbf8c865fd2f

examples/torchrec/main.py

@@ -234,9 +234,11 @@ def init_embdedding_configs(
         EmbeddingBagConfig(
             name=f"t_{feature_name}",
             embedding_dim=embedding_dim,
-            num_embeddings=none_throws(num_embeddings_per_feature)[feature_idx]
-            if num_embeddings is None
-            else num_embeddings,
+            num_embeddings=(
+                none_throws(num_embeddings_per_feature)[feature_idx]
+                if num_embeddings is None
+                else num_embeddings
+            ),
             feature_names=[feature_name],
         )
         for feature_idx, feature_name in enumerate(DEFAULT_CAT_NAMES)

tests/framework/callbacks/test_base_checkpointer.py

@@ -686,33 +686,37 @@ def test_best_checkpoint_no_top_k(self) -> None:
                 bcs.on_train_epoch_end(state, my_train_unit)
                 self.assertEqual(
                     bcs._ckpt_dirpaths,
-                    [
-                        os.path.join(temp_dir, "epoch_1_step_0_train_loss=0.02"),
-                        os.path.join(temp_dir, "epoch_0_step_0_train_loss=0.01"),
-                    ]
-                    if mode == "min"
-                    else [
-                        os.path.join(temp_dir, "epoch_0_step_0_train_loss=0.01"),
-                        os.path.join(temp_dir, "epoch_1_step_0_train_loss=0.02"),
-                    ],
+                    (
+                        [
+                            os.path.join(temp_dir, "epoch_1_step_0_train_loss=0.02"),
+                            os.path.join(temp_dir, "epoch_0_step_0_train_loss=0.01"),
+                        ]
+                        if mode == "min"
+                        else [
+                            os.path.join(temp_dir, "epoch_0_step_0_train_loss=0.01"),
+                            os.path.join(temp_dir, "epoch_1_step_0_train_loss=0.02"),
+                        ]
+                    ),
                 )
 
                 my_train_unit.train_loss = 0.015
                 my_train_unit.train_progress.increment_epoch()
                 bcs.on_train_epoch_end(state, my_train_unit)
                 self.assertEqual(
                     bcs._ckpt_dirpaths,
-                    [
-                        os.path.join(temp_dir, "epoch_1_step_0_train_loss=0.02"),
-                        os.path.join(temp_dir, "epoch_2_step_0_train_loss=0.015"),
-                        os.path.join(temp_dir, "epoch_0_step_0_train_loss=0.01"),
-                    ]
-                    if mode == "min"
-                    else [
-                        os.path.join(temp_dir, "epoch_0_step_0_train_loss=0.01"),
-                        os.path.join(temp_dir, "epoch_2_step_0_train_loss=0.015"),
-                        os.path.join(temp_dir, "epoch_1_step_0_train_loss=0.02"),
-                    ],
+                    (
+                        [
+                            os.path.join(temp_dir, "epoch_1_step_0_train_loss=0.02"),
+                            os.path.join(temp_dir, "epoch_2_step_0_train_loss=0.015"),
+                            os.path.join(temp_dir, "epoch_0_step_0_train_loss=0.01"),
+                        ]
+                        if mode == "min"
+                        else [
+                            os.path.join(temp_dir, "epoch_0_step_0_train_loss=0.01"),
+                            os.path.join(temp_dir, "epoch_2_step_0_train_loss=0.015"),
+                            os.path.join(temp_dir, "epoch_1_step_0_train_loss=0.02"),
+                        ]
+                    ),
                 )
 
     def test_best_checkpoint_top_k(self) -> None:

tests/utils/test_precision.py

@@ -20,7 +20,7 @@
 
 class PrecisionTest(unittest.TestCase):
     def test_convert_precision_str_to_dtype_success(self) -> None:
-        for (precision_str, expected_dtype) in [
+        for precision_str, expected_dtype in [
             ("fp16", torch.float16),
             ("bf16", torch.bfloat16),
             ("fp32", None),

torchtnt/framework/callbacks/base_csv_writer.py

@@ -61,8 +61,7 @@ def get_step_output_rows(
         unit: TPredictUnit,
         # pyre-fixme: Missing parameter annotation [2]
         step_output: Any,
-    ) -> Union[List[str], List[List[str]]]:
-        ...
+    ) -> Union[List[str], List[List[str]]]: ...
 
     def on_predict_start(self, state: State, unit: TPredictUnit) -> None:
         if get_global_rank() == 0:

torchtnt/framework/callbacks/checkpointer_types.py

@@ -7,6 +7,7 @@
 from dataclasses import dataclass
 from typing import Literal, Optional
 
+
 # TODO: eventually support overriding all knobs
 @dataclass
 class KnobOptions:

torchtnt/framework/callbacks/memory_snapshot.py

@@ -15,7 +15,6 @@
 
 
 class MemorySnapshot(Callback):
-
     """
     A callback for memory snapshot collection during training, saving pickle files to the user-specified directory.
     Uses `Memory Snapshots <https://zdevito.github.io/2022/08/16/memory-snapshots.html>`.

torchtnt/utils/data/iterators.py

@@ -332,9 +332,9 @@ def __init__(
                 float(weights[name]) for name in self._iterator_names
             ]
         self._iterator_is_exhausted: List[bool] = [False] * len(self._iterator_names)
-        self.stopping_mechanism: Optional[
-            StoppingMechanism
-        ] = iteration_strategy.stopping_mechanism
+        self.stopping_mechanism: Optional[StoppingMechanism] = (
+            iteration_strategy.stopping_mechanism
+        )
         self.enforce_same_loader_across_ranks: bool = (
             iteration_strategy.enforce_same_loader_across_ranks
         )

torchtnt/utils/loggers/file.py

@@ -89,9 +89,7 @@ def log(self, name: str, data: Scalar, step: int) -> None:
             self._len_before_flush = len(self._log_buffer)
 
     @abstractmethod
-    def flush(self) -> None:
-        ...
+    def flush(self) -> None: ...
 
     @abstractmethod
-    def close(self) -> None:
-        ...
+    def close(self) -> None: ...

torchtnt/utils/memory_snapshot_profiler.py

@@ -57,16 +57,13 @@ def __exit__(
         self.stop()
 
     @abstractmethod
-    def start(self) -> None:
-        ...
+    def start(self) -> None: ...
 
     @abstractmethod
-    def stop(self) -> None:
-        ...
+    def stop(self) -> None: ...
 
     @abstractmethod
-    def step(self) -> None:
-        ...
+    def step(self) -> None: ...
 
 
 class MemorySnapshotProfiler(MemorySnapshotProfilerBase):

torchtnt/utils/stateful.py

@@ -18,11 +18,9 @@
 class Stateful(Protocol):
     """Defines the interface for checkpoint saving and loading."""
 
-    def state_dict(self) -> Dict[str, Any]:
-        ...
+    def state_dict(self) -> Dict[str, Any]: ...
 
-    def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
-        ...
+    def load_state_dict(self, state_dict: Dict[str, Any]) -> None: ...
 
 
 StatefulDict = Dict[str, Stateful]