SparseTensor edge_index not getting batched properly?

[Quoding]

Hi everyone,

I'm running into an issue where I have multiple large graphs that I wish to train on. As such, I have decided to use SparseTensor to improve memory efficiency. However, I am running into an issue, here is a bit of context:

I have a bunch of Data objects in list that I wish to batch, which is symbolically like so:

states_graph = [Data(x=[17217, 2], edge_index=[17217, 17217, nnz=19366912]), Data(x=[17217, 2], edge_index=[17217, 17217, nnz=19366912])]

The edge_index for every Data is a SparseTensor.

I then create a batch from this list:

batch = Batch.from_data_list(states_graph).to(device) # Batch of two graphs

However, the resulting batch does not produce a square sparse matrix for the edge_index:

print(batch_state.edge_index)

# Output:
# SparseTensor(row=tensor([    0,     0,     0,  ..., 17215, 17215, 17215], device='cuda:0'),
#              col=tensor([   32,    33,    34,  ..., 31166, 31167, 31168], device='cuda:0'),
#              size=(17217, 34434), nnz=38733824, density=6.53%)

(i.e. look at the size attribute)
Which seems to show that the columns scale properly with the batch size, but not the rows.

If I run inference on this batch I get

RuntimeError: The size of tensor a (17217) must match the size of tensor b (34434) at non-singleton dimension 0

However, if my batch size is 1, then inference raises no error, which indicates to me that the batching procedure messes with the SparseTensor's dimension.

Is this a bug, or is this expected behavior?

[Quoding]

I fiddled with the torch_sparse code a little bit. It seems that the __cat_dim__ function always returns -1 for the SparseTensor (at least the first in the data list) in the Batch.from_data_list call, in the collate() call within this function.

Setting manually cat_dim = (0, 1) seems to fix the issue in collate.py, L.227

I will investigate deeper, though I would appreciate some input, if you have any.

[Quoding]

Seems to be a PICNIC (Problem in Chair, Not in Computer) error... The way I created my data object stores the adjacency matrix in the edge_index key.

Perhaps it would be a good idea to show how to create such objects in the documentation here?

Looking at the reference for the torch_geometric.data.Data class, it seems like the adj_t attribute is concealed within the **kwargs, which makes in a little bit unclear on how to proceed (e.g. I had assumed only edge_index existed as a valid constructor attribute).

[rusty1s]

Glad you resolved it. You are right that __cat_dim__ currently concatenates edge_index in the last dimension, and adj_t in dimension (0, 1). Let me try to make it more clear in the documentation.