Question about 'index_select' in __lift__ function during graph propagation

[jaes77]

Thanks for your efforts to develop such a great environment.

When launching GCN-like models with edge_index format, I noticed there is a step 'lift' in collect function during propagation.

pytorch_geometric/torch_geometric/nn/conv/message_passing.py

Lines 217 to 219 in 09f25e9

	if isinstance(data, Tensor):
	self.__set_size__(size, dim, data)
	data = self.__lift__(data, edge_index, dim)

pytorch_geometric/torch_geometric/nn/conv/message_passing.py

Lines 186 to 189 in 09f25e9

	def __lift__(self, src, edge_index, dim):
	if isinstance(edge_index, Tensor):
	index = edge_index[dim]
	return src.index_select(self.node_dim, index)

It seems that the step is to arrange input tensor to different form with edge_index. But, I do not catch up what's going on during the step exactly because I'm not familiar with graph models.
What is src.index_select(node_dim, edge_index) exactly and why is the step needed only for 'torch.tensor' type?

I appreciate your support again.

[rusty1s]

The lift call will move node-level tensors to edge-level representations corresponding to source and destination nodes, e.g.

x = ...  # [num_nodes, num_features]
x_src = x_j = x[edge_index[0]]  # [num_edges, num_features]
x_dst = x_i = x[edge_index[1]]  # [num_edges, num_features]

This is needed to compute a message per edge, which we then later aggregate according to destination node indices. Importantly, we can only lift torch.Tensor types since it is the only type that supports this kind of operation.