Make subsample faster in CPU

[fehiepsi]

This PR is motivated by @insperatum in a forum topic and Jake's answer here on the differences between PyTorch CPU and jax implementation.

Benchmarks

import numpyro
from jax import random, jit

size = 100000
subsample_size = 1000

@jit
def subsample_fn(rng_key):
    return numpyro.primitives._subsample_fn(size, subsample_size, rng_key)

key0, key1 = random.PRNGKey(0), random.PRNGKey(1)
x = subsample_fn(key0).copy()
%time x = subsample_fn(key1).copy()

returns

CPU times: user 262 µs, sys: 52 µs, total: 314 µs
Wall time: 253 µs

while in PyTorch %time y = torch.randperm(size)[:subsample_size], it took

CPU times: user 26.5 ms, sys: 1.37 ms, total: 27.9 ms
Wall time: 2.51 ms

and in previous implementation, it took

CPU times: user 60.2 ms, sys: 269 µs, total: 60.4 ms
Wall time: 54.1 ms

The reason for high performance comparing to PyTorch is in PyTorch, we took a permutation of full size first, then collecting a subset. Here, we only take a permutation of size subsample_size. This is observed by @fritzo at this discussion.

[fritzo]

Nice work! We'll have to port this to Pyro if anyone complains about slow subsampling speed there.

[fritzo]

Do you need to ensure double precision so that random.uniform will sample sets larger than 2**24 ~ 16million?

[fehiepsi]

Thanks for reviewing and pointing that issue out, @fritzo! I have switched to randint instead of using uniform. It is a bit slower 451 µs ± 2.06 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each) but is still pretty fast and does not suffer from the precision issue. In JAX, precision is only decided one time, at the beginning of a program, so we can't switch between two modes in the implementation.

[fritzo]

Sorry I lost track of this!