How to convert between dense batches and sparse batches efficiently?

[wzhang2022]

In torch_geometric.utils, there's the to_dense_batch function which allows me to get node features of a batch of graphs as a single tensor of shape (batch_size, n_max, d) along with a boolean mask of shape (batch_size, n_max), where n_max is the size of the largest node. Is there an equivalent function for edge indices, which would output a tensor (batch_size, e_max, 2) along with a mask (batch_size, e_max)? Also, is there a function that would reverse this (i.e. given a dense batch tensor of nodes/edges, output a batch stored as a single graph of batch_size disconnected components the way that PyTorch normally treats batches)?
Additionally, I want to be able to differentiate through this operation as well.

[rusty1s]

There exists the to_dense_adj function that allows to generate dense adjacency matrices of shape [batch_size, max_num_nodes, max_num_nodes], which might be the perfect fit for your use-case.