gigl.src.training.v1.lib.data_loaders.cast_graph_for_training

gigl.src.training.v1.lib.data_loaders.utils.cast_graph_for_training(batch_graph_data: PygGraphData, graph_metadata_pb_wrapper: GraphMetadataPbWrapper, preprocessed_metadata_pb_wrapper: PreprocessedMetadataPbWrapper, batch_type: str, should_register_edge_features: bool | None) → Data | HeteroData

Casts the PygGraphData object into a Data or HeteroData object. Also fills in any missing fields from graph builder with empty tensors in cases where there are no edges for a graph or given edge type. Args:

batch_graph_data (PygGraphData): Coalesced batch graph graph_metadata_pb_wrapper (GraphMetadataPbWrapper): Graph Metadata Pb Wrapper for this training job preprocessed_metadata_pb_wrapper (PreprocessedMetadataPbWrapper): Preprocessed Metadata Pb Wrapper for this training job should_register_edge_features (bool): Whether we should register edge features for the built graph