(g, num_chunks, key, kwargs=None)
| 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 [ |
no test coverage detected