From 7b662e7d59721dc3ff18b43f269fad8a1938a17a Mon Sep 17 00:00:00 2001 From: Joey Yang Date: Mon, 17 Nov 2025 15:42:51 -0800 Subject: [PATCH] Add table_names as a instance variable of SplitTableBatchedEmbeddingBagsCodegen (#5133) Summary: 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 --- .../fbgemm_gpu/split_table_batched_embeddings_ops_training.py | 1 + 1 file changed, 1 insertion(+) diff --git a/fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops_training.py b/fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops_training.py index 0a2a7ab0a1..8b7c5538f3 100644 --- a/fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops_training.py +++ b/fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops_training.py @@ -710,6 +710,7 @@ def __init__( # noqa C901 f"Feature Gates: {[(feature.name, feature.is_enabled()) for feature in FeatureGateName]}" ) + self.table_names: Optional[list[str]] = table_names self.logging_table_name: str = self.get_table_name_for_logging(table_names) self.enable_raw_embedding_streaming: bool = enable_raw_embedding_streaming self.pooling_mode = pooling_mode