Skip to content

Commit 8189ad4

Browse files
Joey Yangmeta-codesync[bot]
authored andcommitted
Add table_names as a instance variable of SplitTableBatchedEmbeddingBagsCodegen (pytorch#5133)
Summary: Pull Request resolved: pytorch#5133 X-link: https://github.com/facebookresearch/FBGEMM/pull/2135 `table_names` is required for raw embedded streaming to support tracking in mpzch modules. Example usage in embedding module to look up raw ids ``` assert hasattr( self.emb_module, "res_params" ), "res_params should exist when raw_id_tracker is enabled" res_params: RESParams = self.emb_module.res_params # pyre-ignore[9] table_names = res_params.table_names raw_ids_dict = raw_id_tracker_wrapper.get_indexed_lookups( table_names, self.emb_module.uuid ) ``` `table_names` only exist when raw embedded streaming is enabled. During testing, we may disable raw embedded streaming (hence no `res_params`) this tight coupling makes testing difficult. This diff adds `table_names` as a instance variable of `SplitTableBatchedEmbeddingBagsCodegen` so we can easily retrieve `table_names` w/o depending on raw embedded streaming. Reviewed By: chouxi Differential Revision: D87089649 fbshipit-source-id: 9f4d28a5d76d7c1fa605639343f1c535a49d5804
1 parent 05d804b commit 8189ad4

File tree

1 file changed

+1
-0
lines changed

1 file changed

+1
-0
lines changed

fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops_training.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -710,6 +710,7 @@ def __init__( # noqa C901
710710
f"Feature Gates: {[(feature.name, feature.is_enabled()) for feature in FeatureGateName]}"
711711
)
712712

713+
self.table_names: Optional[list[str]] = table_names
713714
self.logging_table_name: str = self.get_table_name_for_logging(table_names)
714715
self.enable_raw_embedding_streaming: bool = enable_raw_embedding_streaming
715716
self.pooling_mode = pooling_mode

0 commit comments

Comments
 (0)