RandomLinkSplit behavior with edge_attr

[nilskre]

The RandomLinkSplit splits the edges into training, validation, and test sets. The edge_label_index then contains the edges in the respective train/val/test set.

In the case of edge regression, I want to predict the values at the edge (so the edge_attr). Therefore, I would expect, that the edge_attr is also split into train, validation, and test.

Example shapes:
edge_index:

• data: torch.Size([2, 1134])
• train: torch.Size([2, 795])
• val: torch.Size([2, 795])
• test: torch.Size([2, 908])

edge_label_index:

• train: torch.Size([2, 795])
• val: torch.Size([2, 113])
• test: torch.Size([2, 226])

edge_attr:

• data: torch.Size([1134, 1])
• train: torch.Size([795, 1])
• val: torch.Size([795, 1]) <- Here i would expect similar shapes as the edge_label
• test: torch.Size([908, 1]) <- Here i would expect similar shapes as the edge_label

My code for splitting:

import torch_geometric.transforms as T

split = T.RandomLinkSplit(
    neg_sampling_ratio=0.0,
    # no negative edges for link regression; for link prediction use negative edges
    add_negative_train_samples=False,
    edge_types=edge_types,
    rev_edge_types=reverse_edge_types,
)
train_data, val_data, test_data = split(data)

Thanks in advance!

[nilskre]

Ok, just figured it out. In the default case the edge_attr is not split into train, val, and test. In the default case, an edge_label is created for a binary prediction task.
In the case, that there are already some labels the key parameter of RandomLinkSplit can be used to control the behavior.
So by setting the key to edge_attr the edge_attr gets split up into train, val and test.