Initialize num_chunks for each node/edge. Parameters ---------- g: DGLGraph Graph to be chunked. num_chunks: int Default number of chunks to be applied onto node/edge data. kwargs: dict Key word arguments to specify details for each node/edge data. R
(g, num_chunks, kwargs=None)
| 36 | |
| 37 | |
| 38 | def _initialize_num_chunks(g, num_chunks, kwargs=None): |
| 39 | """Initialize num_chunks for each node/edge. |
| 40 | |
| 41 | Parameters |
| 42 | ---------- |
| 43 | g: DGLGraph |
| 44 | Graph to be chunked. |
| 45 | num_chunks: int |
| 46 | Default number of chunks to be applied onto node/edge data. |
| 47 | kwargs: dict |
| 48 | Key word arguments to specify details for each node/edge data. |
| 49 | |
| 50 | Returns |
| 51 | ------- |
| 52 | num_chunks_data: dict |
| 53 | Detailed number of chunks for each node/edge. |
| 54 | """ |
| 55 | |
| 56 | def _init(g, num_chunks, key, kwargs=None): |
| 57 | chunks_data = kwargs.get(key, None) |
| 58 | is_node = "_node" in key |
| 59 | data_types = g.ntypes if is_node else g.canonical_etypes |
| 60 | if isinstance(chunks_data, int): |
| 61 | chunks_data = {data_type: chunks_data for data_type in data_types} |
| 62 | elif isinstance(chunks_data, dict): |
| 63 | for data_type in data_types: |
| 64 | if data_type not in chunks_data: |
| 65 | chunks_data[data_type] = num_chunks |
| 66 | else: |
| 67 | chunks_data = {data_type: num_chunks for data_type in data_types} |
| 68 | for _, data in chunks_data.items(): |
| 69 | if isinstance(data, dict): |
| 70 | n_chunks = list(data.values()) |
| 71 | else: |
| 72 | n_chunks = [data] |
| 73 | assert all( |
| 74 | isinstance(v, int) for v in n_chunks |
| 75 | ), "num_chunks for each data type should be int." |
| 76 | return chunks_data |
| 77 | |
| 78 | num_chunks_data = {} |
| 79 | for key in [ |
| 80 | "num_chunks_nodes", |
| 81 | "num_chunks_edges", |
| 82 | "num_chunks_node_data", |
| 83 | "num_chunks_edge_data", |
| 84 | ]: |
| 85 | num_chunks_data[key] = _init(g, num_chunks, key, kwargs=kwargs) |
| 86 | return num_chunks_data |
| 87 | |
| 88 | |
| 89 | def _chunk_graph( |