Loads the Cora graph data into MLX array format.
(config)
| 75 | |
| 76 | |
| 77 | def load_data(config): |
| 78 | """Loads the Cora graph data into MLX array format.""" |
| 79 | print("Loading Cora dataset...") |
| 80 | |
| 81 | # Download dataset files |
| 82 | download_cora() |
| 83 | |
| 84 | # Graph nodes |
| 85 | raw_nodes_data = np.genfromtxt(config.nodes_path, dtype="str") |
| 86 | raw_node_ids = raw_nodes_data[:, 0].astype( |
| 87 | "int32" |
| 88 | ) # unique identifier of each node |
| 89 | raw_node_labels = raw_nodes_data[:, -1] |
| 90 | labels_enumerated = enumerate_labels(raw_node_labels) # target labels as integers |
| 91 | node_features = sparse.csr_matrix(raw_nodes_data[:, 1:-1], dtype="float32") |
| 92 | |
| 93 | # Edges |
| 94 | ids_ordered = {raw_id: order for order, raw_id in enumerate(raw_node_ids)} |
| 95 | raw_edges_data = np.genfromtxt(config.edges_path, dtype="int32") |
| 96 | edges_ordered = np.array( |
| 97 | list(map(ids_ordered.get, raw_edges_data.flatten())), dtype="int32" |
| 98 | ).reshape(raw_edges_data.shape) |
| 99 | |
| 100 | # Adjacency matrix |
| 101 | adj = sparse.coo_matrix( |
| 102 | (np.ones(edges_ordered.shape[0]), (edges_ordered[:, 0], edges_ordered[:, 1])), |
| 103 | shape=(labels_enumerated.shape[0], labels_enumerated.shape[0]), |
| 104 | dtype=np.float32, |
| 105 | ) |
| 106 | |
| 107 | # Make the adjacency matrix symmetric |
| 108 | adj = adj + adj.T.multiply(adj.T > adj) |
| 109 | adj = normalize_adjacency(adj) |
| 110 | |
| 111 | # Convert to mlx array |
| 112 | features = mx.array(node_features.toarray(), mx.float32) |
| 113 | labels = mx.array(labels_enumerated, mx.int32) |
| 114 | adj = mx.array(adj.toarray()) |
| 115 | |
| 116 | print("Dataset loaded.") |
| 117 | return features, labels, adj |
no test coverage detected