Repository: google/jax

Issue:
Add jax.random.choice
It would be nice to have a random weighted choice in jax. We need this to sample bitstrings from a quantum wave function.


Discussion:
I ran into the same issue, I got around it by implementing the function described in this [blog post](https://timvieira.github.io/blog/post/2019/09/16/algorithms-for-sampling-without-replacement/). Ideally `jax.random.choice` would be implemented, but I suppose it's tricky to get it to match the API for `np.random.choice`. Perhaps `jax.random` could have a simpler sampling without replacement function?

Edit: Seems like #2066 discusses similar issues in how to match the numpy version.
This is a significant missing feature. I'm curious @Thenerdstation if you also need sampling without replacement or if you're happy with a single sample/sampling with replacement (both much easier).
Sampling with replacement is what we need for our use case
seconded

should this take a probability vector `p`? `categorical` takes a `logits` vector. (was this for performance reasons?)

guessing api would be something like 
```
jax.random.choice(key, x, size=None, replace=True, p=None, axis=0)
```
potentially relevant: `permutation` was requested https://github.com/google/jax/issues/1526 with PR pending https://github.com/google/jax/pull/1568

I also have been in need of this several times.
Both the with and without replacement case.

Resolving patch:
diff --git a/jax/random.py b/jax/random.py
--- a/jax/random.py
+++ b/jax/random.py
@@ -530,6 +530,63 @@ def _shuffle(key, x, axis):
   return x
 
 
+def choice(key, a, shape=(), replace=True, p=None):
+  """Generates a random sample from a given 1-D array.
+
+  Args:
+    key: a PRNGKey used as the random key.
+    a : 1D array or int. If an ndarray, a random sample is generated from
+      its elements. If an int, the random sample is generated as if a were
+      arange(a).
+    shape : tuple of ints, optional. Output shape.  If the given shape is,
+      e.g., ``(m, n)``, then ``m * n`` samples are drawn.  Default is (),
+      in which case a single value is returned.
+    replace : boolean.  Whether the sample is with or without replacement.
+      default is True.
+    p : 1-D array-like, The probabilities associated with each entry in a.
+      If not given the sample assumes a uniform distribution over all
+      entries in a.
+
+  Returns:
+    An array of shape `shape` containing samples from `a`.
+  """
+  a = jnp.asarray(a)
+  if a.ndim not in [0, 1]:
+    raise ValueError("a must be an integer or 1-dimensional")
+  n_inputs = int(a) if a.ndim == 0 else len(a)
+  n_draws = np.prod(shape).astype(int)
+  if n_draws == 0:
+    return jnp.zeros(shape, dtype=a.dtype)
+  if n_inputs <= 0:
+    raise ValueError("a must be greater than 0 unless no samples are taken")
+  if not replace and n_draws > n_inputs:
+    raise ValueError("Cannot take a larger sample than population when 'replace=False'")
+
+  if p is None:
+    if replace:
+      ind = randint(key, shape, 0, n_inputs)
+      result = ind if a.ndim == 0 else a[ind]
+    else:
+      result = permutation(key, a)[:n_draws]
+  else:
+    p = jnp.asarray(p)
+    if p.shape != (n_inputs,):
+      raise ValueError("p must be None or match the shape of a")
+    if jnp.any(p < 0):
+      raise ValueError("entries of p must be non-negative.")
+    if replace:
+      p_cuml = jnp.cumsum(p)
+      r = p_cuml[-1] * (1 - uniform(key, shape))
+      ind = jnp.searchsorted(p_cuml, r)
+      result = ind if a.ndim == 0 else a[ind]
+    else:
+      # Gumbel top-k trick: https://timvieira.github.io/blog/post/2019/09/16/algorithms-for-sampling-without-replacement/
+      g = -gumbel(key, (n_inputs,)) - jnp.log(p)
+      ind = jnp.argsort(g)[:n_draws]
+      result = ind if a.ndim == 0 else a[ind]
+  return result.reshape(shape)
+
+
 def normal(key: jnp.ndarray,
            shape: Sequence[int] = (),
            dtype: np.dtype = np.float64) -> jnp.ndarray:


Test patch:
diff --git a/tests/random_test.py b/tests/random_test.py
--- a/tests/random_test.py
+++ b/tests/random_test.py
@@ -233,6 +233,34 @@ def testShuffle(self, dtype):
     self.assertFalse(np.all(perm1 == x))  # seems unlikely!
     self.assertAllClose(np.sort(perm1), x, check_dtypes=False)
 
+  @parameterized.named_parameters(jtu.cases_from_list(
+      {"testcase_name": "_{}_shape={}_replace={}_weighted={}_array_input={}".format(
+          np.dtype(dtype).name, shape, replace, weighted, array_input),
+        "dtype": np.dtype(dtype).name, "shape": shape, "replace": replace,
+        "weighted": weighted, "array_input": array_input}
+      for dtype in [np.float32, np.float64, np.int32, np.int64]
+      for shape in [(), (5,), (4, 5)]
+      for replace in [True, False]
+      for weighted in [True, False]
+      for array_input in [True, False]))
+  def testChoice(self, dtype, shape, replace, weighted, array_input):
+    N = 100
+    key = random.PRNGKey(0)
+    x = N if not array_input else jnp.arange(N, dtype=dtype)
+    p = None if not weighted else jnp.arange(N)
+    rand = lambda key: random.choice(key, x, shape, p=p, replace=replace)
+    crand = api.jit(rand)
+
+    sample1 = rand(key)
+    sample2 = crand(key)
+
+    self.assertEqual(shape, sample1.shape)
+    if array_input:
+      self.assertEqual(x.dtype, sample1.dtype)
+    if not replace:
+      assert len(np.unique(sample1)) == len(np.ravel(sample1))
+    self.assertAllClose(sample1, sample2)
+
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_{}".format(jtu.format_shape_dtype_string(shape, dtype)),
        "dtype": np.dtype(dtype).name, "shape": shape}