Monday 2 August 2021

Efficient multinomial sampling for sparse array/tensor in python

I have a sparse array/tensor like below.

import torch
from torch_sparse import SparseTensor


row = torch.tensor([0, 0, 0, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 3])
col = torch.tensor([1, 2, 3, 0, 2, 0, 1, 4, 5, 0, 2, 5, 2, 4])
value = torch.rand([14])
adj_t = SparseTensor(row=row, col=col, value=value, sparse_sizes=(4, 9))

I want to sample n_samples column index's with or without replacement. I can do this by first converting adj_t to dense and then using torch.multinomial or similarly with numpy.random.choice.

sample = torch.multinomial(adj_t.to_dense(), num_samples=2, replacement=True)

But converting the sparse array to dense and the torch.multinomial is not very efficient. Is there a sparse version of torch.multinomial. If no, how would one go about implementing it



from Efficient multinomial sampling for sparse array/tensor in python

No comments:

Post a Comment