[Automated Commit] Format Codebase

Author: github-actions[bot]

recommendation/dlrm_v3/accuracy.py

@@ -67,9 +67,10 @@ def main() -> None:
         num_candidates = data[-1].astype(int)
         assert len(data) == 1 + num_candidates * 3
         mt_target_preds = torch.from_numpy(data[0:num_candidates])
-        mt_target_labels = torch.from_numpy(data[num_candidates : num_candidates * 2])
+        mt_target_labels = torch.from_numpy(
+            data[num_candidates: num_candidates * 2])
         mt_target_weights = torch.from_numpy(
-            data[num_candidates * 2 : num_candidates * 3]
+            data[num_candidates * 2: num_candidates * 3]
         )
         num_candidates = torch.tensor([num_candidates])
         metrics.update(

recommendation/dlrm_v3/checkpoint.py

@@ -46,7 +46,8 @@ class SparseState(Stateful):
         sparse_tensor_keys: Set of keys identifying sparse tensors in the model's state dict.
     """
 
-    def __init__(self, model: torch.nn.Module, sparse_tensor_keys: Set[str]) -> None:
+    def __init__(self, model: torch.nn.Module,
+                 sparse_tensor_keys: Set[str]) -> None:
         self.model = model
         self.sparse_tensor_keys = sparse_tensor_keys
 
@@ -62,17 +63,23 @@ def state_dict(self) -> Dict[str, torch.Tensor]:
         return out_dict
 
     def load_state_dict(self, state_dict: Dict[str, torch.Tensor]) -> None:
-        incompatible_keys = self.model.load_state_dict(state_dict, strict=False)
+        incompatible_keys = self.model.load_state_dict(
+            state_dict, strict=False)
         assert not incompatible_keys.unexpected_keys
 
 
 def is_sparse_key(k: str, v: torch.Tensor) -> bool:
     return isinstance(v, ShardedTensor) or "embedding_collection" in k
 
 
-def load_dense_state_dict(model: torch.nn.Module, state_dict: Dict[str, Any]) -> None:
+def load_dense_state_dict(model: torch.nn.Module,
+                          state_dict: Dict[str, Any]) -> None:
     own_state = model.state_dict()
-    own_state_dense_keys = {k for k, v in own_state.items() if not is_sparse_key(k, v)}
+    own_state_dense_keys = {
+        k for k,
+        v in own_state.items() if not is_sparse_key(
+            k,
+            v)}
     state_dict_dense_keys = {
         k for k, v in state_dict.items() if not is_sparse_key(k, v)
     }
@@ -156,7 +163,8 @@ def save_dmp_checkpoint(
     sparse_dict = {"sparse_dict": SparseState(model, sparse_tensor_keys)}
     torch.distributed.checkpoint.save(
         sparse_dict,
-        storage_writer=torch.distributed.checkpoint.FileSystemWriter(sparse_path),
+        storage_writer=torch.distributed.checkpoint.FileSystemWriter(
+            sparse_path),
     )
     torch.distributed.barrier()
     print("checkpoint successfully saved")
@@ -178,7 +186,8 @@ def load_sparse_checkpoint(
     gc.collect()
     torch.distributed.checkpoint.load(
         sparse_dict,
-        storage_reader=torch.distributed.checkpoint.FileSystemReader(sparse_path),
+        storage_reader=torch.distributed.checkpoint.FileSystemReader(
+            sparse_path),
     )
     gc.collect()
     print("sparse checkpoint successfully loaded")

recommendation/dlrm_v3/configs.py

@@ -114,7 +114,8 @@ def get_hstu_configs(dataset: str = "debug") -> DlrmHSTUConfig:
     return hstu_config
 
 
-def get_embedding_table_config(dataset: str = "debug") -> Dict[str, EmbeddingConfig]:
+def get_embedding_table_config(
+        dataset: str = "debug") -> Dict[str, EmbeddingConfig]:
     """
     Create and return embedding table configurations.
 

recommendation/dlrm_v3/data_producer.py

@@ -90,7 +90,8 @@ def enqueue(
         """
         with torch.profiler.record_function("data batching"):
             t0_batching: float = time.time()
-            samples: Union[Samples, List[Samples]] = self.ds.get_samples(content_ids)
+            samples: Union[Samples, List[Samples]
+                           ] = self.ds.get_samples(content_ids)
             dt_batching: float = time.time() - t0_batching
             if isinstance(samples, Samples):
                 query = QueryItem(
@@ -106,7 +107,7 @@ def enqueue(
                 for sample in samples:
                     batch_size: int = sample.batch_size()
                     query = QueryItem(
-                        query_ids=query_ids[start_idx : start_idx + batch_size],
+                        query_ids=query_ids[start_idx: start_idx + batch_size],
                         samples=sample,
                         start=t0,
                         dt_queue=dt_queue,
@@ -148,7 +149,9 @@ def __init__(
         )
         self.workers: List[threading.Thread] = []
         for _ in range(self.threads):
-            worker = threading.Thread(target=self.handle_tasks, args=(self.tasks,))
+            worker = threading.Thread(
+                target=self.handle_tasks, args=(
+                    self.tasks,))
             worker.daemon = True
             self.workers.append(worker)
             worker.start()
@@ -172,7 +175,8 @@ def handle_tasks(
                 break
             query_ids, content_ids, t0, dt_queue = query_and_content_ids
             t0_batching: float = time.time()
-            samples: Union[Samples, List[Samples]] = self.ds.get_samples(content_ids)
+            samples: Union[Samples, List[Samples]
+                           ] = self.ds.get_samples(content_ids)
             dt_batching: float = time.time() - t0_batching
             if isinstance(samples, Samples):
                 qitem = QueryItem(
@@ -189,7 +193,7 @@ def handle_tasks(
                 for sample in samples:
                     batch_size: int = sample.batch_size()
                     qitem = QueryItem(
-                        query_ids=query_ids[start_idx : start_idx + batch_size],
+                        query_ids=query_ids[start_idx: start_idx + batch_size],
                         samples=sample,
                         start=t0,
                         dt_queue=dt_queue,

recommendation/dlrm_v3/datasets/dataset.py

@@ -204,11 +204,13 @@ def kjt_batch_func(
     bs_offset = torch.ops.fbgemm.asynchronous_complete_cumsum(
         torch.tensor(bs_list)
     ).int()
-    batched_offset = torch.ops.fbgemm.asynchronous_complete_cumsum(batched_length)
+    batched_offset = torch.ops.fbgemm.asynchronous_complete_cumsum(
+        batched_length)
     reorder_length = torch.ops.fbgemm.reorder_batched_ad_lengths(
         batched_length, bs_offset, bs
     )
-    reorder_offsets = torch.ops.fbgemm.asynchronous_complete_cumsum(reorder_length)
+    reorder_offsets = torch.ops.fbgemm.asynchronous_complete_cumsum(
+        reorder_length)
     reorder_indices = torch.ops.fbgemm.reorder_batched_ad_indices(
         batched_offset, batched_indices, reorder_offsets, bs_offset, bs
     )
@@ -345,7 +347,8 @@ def __init__(
         self.num_aggregated_samples = num_aggregated_samples
         self.items_in_memory = {}
 
-    def get_sample(self, id: int) -> Tuple[KeyedJaggedTensor, KeyedJaggedTensor]:
+    def get_sample(
+            self, id: int) -> Tuple[KeyedJaggedTensor, KeyedJaggedTensor]:
         """
         Get a sample by ID from in-memory storage.
 

recommendation/dlrm_v3/datasets/synthetic_streaming.py

@@ -146,7 +146,8 @@ def load_query_samples(self, sample_list: List[int]) -> None:
     def unload_query_samples(self, sample_list: List[int]) -> None:
         self.items_in_memory = {}
 
-    def get_sample(self, id: int) -> Tuple[KeyedJaggedTensor, KeyedJaggedTensor]:
+    def get_sample(
+            self, id: int) -> Tuple[KeyedJaggedTensor, KeyedJaggedTensor]:
         return self.items_in_memory[self.ts][id]
 
     def get_sample_with_ts(
@@ -192,7 +193,8 @@ def _process_line(self, line: str, user_id: int) -> pd.Series:
         reader = csv.reader([line])
         parsed_line = next(reader)
         # total ts + one more eval ts + one base ts so that uih won't be zero
-        # for each ts, ordered as candidate_ids, candidate_ratings, uih_ids, uih_ratings
+        # for each ts, ordered as candidate_ids, candidate_ratings, uih_ids,
+        # uih_ratings
         assert len(parsed_line) == 4 * (self.total_ts + 2)
         uih_item_ids_list = []
         uih_ratings_list = []
@@ -290,7 +292,8 @@ def set_ts(self, ts: int) -> None:
             assert len(row) == 1
             requests = json_loads(row[0])
             self.requests = requests
-            logger.warning(f"DLRMv3SyntheticStreamingDataset: ts={ts} requests loaded")
+            logger.warning(
+                f"DLRMv3SyntheticStreamingDataset: ts={ts} requests loaded")
         assert self.ts_to_users_cumsum[self.ts][-1] == len(self.requests)
         logger.warning(
             f"DLRMv3SyntheticStreamingDataset: ts={ts} users_cumsum={self.ts_to_users_cumsum[self.ts]}"
@@ -336,7 +339,8 @@ def load_item(
         timestamps_uih = maybe_truncate_seq(timestamps_uih, self._max_uih_len)
         ids_candidates = maybe_truncate_seq(ids_candidates, max_num_candidates)
         num_candidates = len(ids_candidates)
-        ratings_candidates = maybe_truncate_seq(ratings_candidates, max_num_candidates)
+        ratings_candidates = maybe_truncate_seq(
+            ratings_candidates, max_num_candidates)
         action_weights_uih = [
             self.action_weights[int(rating) - 1] for rating in ratings_uih
         ]
@@ -366,7 +370,8 @@ def load_item(
             [
                 uih_seq_len
                 for _ in range(
-                    len(self._uih_keys) - len(self._contextual_feature_to_max_length)
+                    len(self._uih_keys) -
+                    len(self._contextual_feature_to_max_length)
                 )
             ]
         )
@@ -380,7 +385,8 @@ def load_item(
             values=torch.tensor(uih_kjt_values).long(),
         )
 
-        candidates_kjt_lengths = num_candidates * torch.ones(len(self._candidates_keys))
+        candidates_kjt_lengths = num_candidates * \
+            torch.ones(len(self._candidates_keys))
         item_candidate_category_ids = [
             id // self.items_per_category for id in ids_candidates
         ]

recommendation/dlrm_v3/datasets/utils.py

@@ -45,7 +45,7 @@ def json_loads(
         y = json.loads(x)
     else:
         y = x
-    y_list = [y] if type(y) == int else list(y)
+    y_list = [y] if isinstance(y, int) else list(y)
     return y_list
 
 
@@ -72,7 +72,7 @@ def separate_uih_candidates(
         y = json.loads(x)
     else:
         y = x
-    y_list = [y] if type(y) == int else list(y)
+    y_list = [y] if isinstance(y, int) else list(y)
     candidates, uih = (
         y_list[-candidates_max_seq_len:],
         y_list[:-candidates_max_seq_len],

recommendation/dlrm_v3/generative_recommenders/common.py

@@ -188,7 +188,8 @@ def generate_sparse_seq_len(
     if sparsity == 0.0:
         return torch.zeros(size=(size,), device=device, dtype=torch.int)
     elif sparsity == 1.0:
-        return torch.ones(size=(size,), device=device, dtype=torch.int) * max_seq_len
+        return torch.ones(size=(size,), device=device,
+                          dtype=torch.int) * max_seq_len
     elif sparsity >= 0.5:
         min_seq_len: int = int((2 * sparsity - 1.0) * max_seq_len)
         return torch.randint(
@@ -265,10 +266,12 @@ def switch_to_contiguous_if_needed(x: torch.Tensor) -> torch.Tensor:
 def prev_power_of_2(x: int) -> int:
     if torch.compiler.is_compiling():
         # Re-write to make Dynamo happy
-        x_tensor = torch.scalar_tensor(x, dtype=torch.int64)  # type: ignore[arg-type]
+        x_tensor = torch.scalar_tensor(
+            x, dtype=torch.int64)  # type: ignore[arg-type]
         x_tensor_orig = x_tensor.clone()
         out = triton.next_power_of_2(x_tensor)  # type: ignore[arg-type]
-        return int(torch.where(torch.lt(x_tensor_orig, out), out // 2, out).item())  # type: ignore[return-value]
+        return int(torch.where(torch.lt(x_tensor_orig, out), out //
+                   2, out).item())  # type: ignore[return-value]
     else:
         out = triton.next_power_of_2(x)
         return out // 2 if out > x else out
@@ -340,7 +343,9 @@ def _generate_fine_grained_buckets() -> List[int]:
 def _fine_grained_bucket_size(x: int) -> int:
     if torch.compiler.is_compiling():
         x_tensor = torch.scalar_tensor(x, dtype=torch.int64)
-        buckets = torch.tensor(_generate_fine_grained_buckets(), dtype=torch.int64)
+        buckets = torch.tensor(
+            _generate_fine_grained_buckets(),
+            dtype=torch.int64)
 
         mask = buckets >= x_tensor
         valid_buckets = torch.where(
@@ -361,7 +366,8 @@ def _fine_grained_bucket_size(x: int) -> int:
 
 
 @torch.fx.wrap
-def fx_unwrap_optional_tensor(optional: Optional[torch.Tensor]) -> torch.Tensor:
+def fx_unwrap_optional_tensor(
+        optional: Optional[torch.Tensor]) -> torch.Tensor:
     assert optional is not None, "Expected optional to be non-None Tensor"
     return optional
 

recommendation/dlrm_v3/generative_recommenders/modules/action_encoder.py

@@ -85,7 +85,8 @@ def forward(
             watchtimes = seq_payloads[self._watchtime_feature_name]
             for threshold, weight in self._watchtime_to_action_thresholds_and_weights:
                 seq_actions = torch.bitwise_or(
-                    seq_actions, (watchtimes >= threshold).to(torch.int64) * weight
+                    seq_actions, (watchtimes >= threshold).to(
+                        torch.int64) * weight
                 )
         exploded_actions = (
             torch.bitwise_and(
@@ -94,7 +95,8 @@ def forward(
             > 0
         )
         action_embeddings = (
-            exploded_actions.unsqueeze(-1) * self._action_embedding_table.unsqueeze(0)
+            exploded_actions.unsqueeze(-1) *
+            self._action_embedding_table.unsqueeze(0)
         ).view(-1, self._num_action_types * self._action_embedding_dim)
         total_targets: int = seq_embeddings.size(0) - action_embeddings.size(0)
         action_embeddings = concat_2D_jagged(

recommendation/dlrm_v3/generative_recommenders/modules/content_encoder.py

@@ -79,7 +79,8 @@ def forward(
         if self._target_enrich_dummy_embeddings:
             total_seq_len: int = seq_embeddings.size(0)
             for name, param in self._target_enrich_dummy_embeddings.items():
-                enrich_embeddings_target = seq_payloads[name].to(seq_embeddings.dtype)
+                enrich_embeddings_target = seq_payloads[name].to(
+                    seq_embeddings.dtype)
                 total_targets: int = enrich_embeddings_target.size(0)
                 total_uih_len: int = total_seq_len - total_targets
                 enrich_embeddings_uih = param.tile(total_uih_len, 1).to(