Any efficient way to merge edge features from one graph to another?

[SwaySZ]

Hi,
I have two graphs A and B. There might be some different edges between them. I wanna merge the edge features from A to B but drop out the edges that are not in B. So, I need to first check if an edge of A is in B. One possible way is as follows:

#given features x_A and x_B
x_A_tmp = torch.zero_like(x_B)
k = 0
for edge_index in B.edge_index:
    if edge_index in A.edge_index:
        x_A_tmp[k] = x_A[k]
    else:
        x_A_tmp[k] = 0
    k++

I think the above code is very inefficient, so look for a better one. Thank you.

Sway

[rusty1s]

In PyG, we usually do this by converting edge indices into a single 1-dimensional index tensor, and then performing the filtering via torch.isin, e.g.:

index1 = edge_index_A[0] * num_nodes + edge_index_A[1]
index2 = edge_index_B[0] * num_nodes + edge_index_B[1]
mask = torch.isin(index1, index2)

[SwaySZ]

Dear Dr. Fey,

Thanks for your solution. I appreciate it. The following is my code:

mask1 = torch.isin(index1, index2, assume_unique = True)
mask2 = torch.isin(index2, index1, assume_unique = True)
B_edge_features[mask2] = A_edge_features[mask1]

However, I have to do twice isin check. Is there any possible improvement? Thank you!

Sway

[rusty1s]

I don't think you can do it without two isin checks TBH.