tensorflow biased sample from sparse matrix

Question

I have a sparse adjacency matrix (around 1.5mln X 1/5 mln symmetrical) with 50% rows < 5 elements, whose sparse representation fits in the gpu entirely (around 3GB). I would like to sample up to n (32) elements from interested rows of a batch (size 1024). And this should be as quick as possible.

I also aim to

• Biased sampling in function of the logarithmic strength of the values
• Sampling at multiple (example 2) hops

So far I was using dgl to do that but it is incredibly slow and add a tremendous latency to the current model.

Can you please suggest other ways for doing this (hopefully another graph library which translate to tensorflow well) or would you suggest implement the sampling with the sparse tensor data structure?

Thank you very much