I’m using PyTorch Geometric, and I’m interested in converting a batch of graphs into separate batches of node features and edge features. In particular, I’m interested in being able to separate the nodes along the dimension by which graph they belong to, but I’m only given the batches in this format.
Concretely, I have tensors batch = Tensor([0, 0, 0, 0, 1, 1, 2]), nodes = Tensor([1, 2, 3, 4, 5, 6, 7]), and I want to convert it to a tensor that looks something like nodes = Tensor([[1,2,3,4], [5,6, pad, pad], [7, pad, pad, pad]) where I use pad to denote some padding tokens. In this case, I have turned a tensor of shape (7,) (number of nodes) into a tensor of shape (3, 4) (number of graphs in batch, maximum number of nodes per graph).