Make subsample faster in CPU

numpyro/primitives.py

@@ -6,7 +6,8 @@
 import functools
 import warnings
 
-from jax import lax, random
+from jax import lax, ops, random
+from jax.lib import xla_bridge
 import jax.numpy as jnp
 
 import numpyro
@@ -235,7 +236,21 @@ def module(name, nn, input_shape=None):
 
 def _subsample_fn(size, subsample_size, rng_key=None):
     assert rng_key is not None, "Missing random key to generate subsample indices."
-    return random.permutation(rng_key, size)[:subsample_size]
+    if xla_bridge.get_backend().platform == 'cpu':
+        # ref: https://en.wikipedia.org/wiki/Fisher%E2%80%93Yates_shuffle#The_modern_algorithm
+        rng_keys = random.split(rng_key, subsample_size)
+
+        def body_fn(val, idx):
+            i_p1 = size - idx
+            i = i_p1 - 1
+            j = random.randint(rng_keys[idx], (), 0, i_p1)
+            val = ops.index_update(val, ops.index[[i, j], ], val[ops.index[[j, i], ]])
+            return val, None
+
+        val, _ = lax.scan(body_fn, jnp.arange(size), jnp.arange(subsample_size))
+        return val[-subsample_size:]
+    else:
+        return random.choice(rng_key, size, (subsample_size,), replace=False)
 
 
 class plate(Messenger):

test/test_handlers.py

@@ -612,3 +612,22 @@ def model(x=None):
     z0 = handlers.seed(model, 0)()
     assert (z[1:] != z0).all()
     assert (z[0] == z0).all()
+
+
+def test_subsample_fn():
+    size = 20
+    subsample_size = 11
+    num_samples = 1000000
+
+    @jit
+    def subsample_fn(rng_key):
+        return numpyro.primitives._subsample_fn(size, subsample_size, rng_key)
+
+    rng_keys = random.split(random.PRNGKey(0), num_samples)
+    subsamples = vmap(subsample_fn)(rng_keys)
+    for k in range(1, 11):
+        i = random.randint(random.PRNGKey(k), (), 0, size)
+        assert_allclose(jnp.mean(subsamples == i, axis=0), jnp.full(subsample_size, 1 / size), atol=1e-3)
+
+        # test that values are not duplicated
+        assert len(set(subsamples[k])) == subsample_size