Inverse of to_dense_batch?

[Boltzmachine]

Is there a function which is the inverse of to_dense_batch? i.e. takes in a tensor of size [bs, Nmax, dim] and a mask of size [bs, Nmax], and returns the tensor of size [N, dim] and its batch indicator batch of size [N].
I notice there's a function called dense_to_sparse, but it only deals with adj.

[rusty1s]

You can simply do

x, mask = to_dense_batch(x, batch)
x = x[mask]

[Boltzmachine]

Oh great thanks. But Is there a function that converts a batch of adj and mask to edge_index?

[rusty1s]

See https://pytorch-geometric.readthedocs.io/en/latest/modules/utils.html#torch_geometric.utils.dense_to_sparse