Graph Readout

Pool node embeddings into one vector per graph, for graph-level tasks such as predicting a property of a molecule.

Mini-batches of graphs are usually stored as one big disconnected graph. A graph_ids array says which graph each node belongs to. After the message passing layers, a readout combines the node rows of each graph into a single row:

• "sum": add the node embeddings
• "mean": average them
• "max": take the maximum of each feature

The readout must not depend on the order of the nodes, because a graph has no natural node order.

Function signature

def graph_readout(H, graph_ids, num_graphs, mode="mean"):

Example

Input

H = np.array([[1.0, 2.0],
              [3.0, 4.0],
              [5.0, 6.0]])
graph_readout(H, np.array([0, 0, 1]), num_graphs=2, mode="mean")

Output

array([[2., 3.],
       [5., 6.]])

Reasoning

Nodes 0 and 1 belong to graph 0, so its row is the mean of [1, 2] and [3, 4]. Node 2 is the only node of graph 1.

Constraints and notes

• graph_ids holds integers in [0, num_graphs), in any order.
• A graph with no nodes gets a row of zeros in every mode.
• Return an array of shape (num_graphs, d).