Repository: google/jax

Issue:
Implement stratified resampling np.digitize in XLA
Hi JAX team,

I'm interested in implementing efficient stratified resampling in JAX. This is a technique for variance reduction in sequential monte carlo algorithms (see section 2.2.1 in Naesseth, Lindsten, and Schön "Elements of Sequential Monte Carlo").

The algorithm provided looks like this

```python
# w.shape = (N, ) 
# x.shape = (N, )
# assert np.sum(w) = 1
def stratified_resampling(key, w, x) :
  N = w.shape[0]
  u = (np.arange(N) + jax.random.uniform(key, (N, )) / N
  bins = np.cumsum(w)
  return x[np.digitize(u, bins)]
```

although I'm also considering a version that doesn't actually do the shuffling in the last step (saving it for a "decoding" process later on which doesn't particularly need to be fast.

The immediate issue is that `np.digitize` is not implemented in JAX. As it were, I don't actually need the gradient for it, but as it's in the hot path and defeats the tracer it's causing some real pain. Even more ideally, I'd like to have something like the interface to `jax.random.categorical` where I can just use unnormalized logits.

How can I help implement at least tracing through `np.digitize`? Alternatively, is there a clever way of doing stratified resampling without using `np.digitize`?

Thanks


Discussion:
@tel I had the same issue. `digitize` relies on `searchsorted`. The latter is not implemented in JAX either, but you can sort of hack it if you don't mind a few extra allocations and sorts, see https://github.com/google/jax/issues/2080#issuecomment-616234407. 
@jakevdp interested?

Resolving patch:
diff --git a/jax/numpy/__init__.py b/jax/numpy/__init__.py
--- a/jax/numpy/__init__.py
+++ b/jax/numpy/__init__.py
@@ -28,10 +28,10 @@
     complex128, complex64, complex_, complexfloating, concatenate, conj,
     conjugate, convolve, copysign, corrcoef, correlate, cos, cosh,
     count_nonzero, cov, cross, csingle, cumprod, cumproduct, cumsum, deg2rad,
-    degrees, diag, diag_indices, diagonal, diff, divide, divmod, dot, double,
-    dsplit, dstack, dtype, e, ediff1d, einsum, einsum_path, empty, empty_like,
-    equal, euler_gamma, exp, exp2, expand_dims, expm1, eye, fabs, finfo, fix,
-    flexible, flip, fliplr, flipud, float16, float32, float64, float_,
+    degrees, diag, diag_indices, diagonal, diff, digitize, divide, divmod, dot,
+    double, dsplit, dstack, dtype, e, ediff1d, einsum, einsum_path, empty,
+    empty_like, equal, euler_gamma, exp, exp2, expand_dims, expm1, eye, fabs,
+    finfo, fix, flexible, flip, fliplr, flipud, float16, float32, float64, float_,
     float_power, floating, floor, floor_divide, fmax, fmin, fmod, frexp, full,
     full_like, function, gcd, geomspace, gradient, greater, greater_equal,
     hamming, hanning, heaviside, hsplit, hstack, hypot, identity, iinfo, imag,
diff --git a/jax/numpy/lax_numpy.py b/jax/numpy/lax_numpy.py
--- a/jax/numpy/lax_numpy.py
+++ b/jax/numpy/lax_numpy.py
@@ -3823,6 +3823,18 @@ def searchsorted(a, v, side='left', sorter=None):
   return _searchsorted(a, v, side)
 
 
+@_wraps(onp.digitize)
+def digitize(x, bins, right=False):
+  if len(bins) == 0:
+    return zeros(x, dtype=int32)
+  side = 'right' if not right else 'left'
+  return where(
+    bins[-1] >= bins[0],
+    searchsorted(bins, x, side=side),
+    len(bins) - searchsorted(bins[::-1], x, side=side)
+  )
+
+
 @_wraps(onp.percentile)
 def percentile(a, q, axis=None, out=None, overwrite_input=False,
                interpolation="linear", keepdims=False):


Test patch:
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -1431,6 +1431,28 @@ def testSearchsorted(self, ashape, vshape, side, dtype, rng_factory):
     self._CheckAgainstNumpy(onp_fun, jnp_fun, args_maker, check_dtypes=True)
     self._CompileAndCheck(jnp_fun, args_maker, check_dtypes=True)
 
+  @parameterized.named_parameters(jtu.cases_from_list(
+    {"testcase_name": "_x={}_bins={}_right={}_reverse={}".format(
+      jtu.format_shape_dtype_string(xshape, dtype),
+      jtu.format_shape_dtype_string(binshape, dtype),
+      right, reverse), "xshape": xshape, "binshape": binshape,
+      "right": right, "reverse": reverse, "dtype": dtype, "rng_factory": rng_factory}
+    for xshape in [(20,), (5, 4)]
+    for binshape in [(1,), (5,)]
+    for right in [True, False]
+    for reverse in [True, False]
+    for dtype in default_dtypes
+    for rng_factory in [jtu.rand_default]
+  ))
+  def testDigitize(self, xshape, binshape, right, reverse, dtype, rng_factory):
+    order = jax.ops.index[::-1] if reverse else jax.ops.index[:]
+    rng = rng_factory(self.rng())
+    args_maker = lambda: [rng(xshape, dtype), jnp.sort(rng(binshape, dtype))[order]]
+    onp_fun = lambda x, bins: onp.digitize(x, bins, right=right)
+    jnp_fun = lambda x, bins: jnp.digitize(x, bins, right=right)
+    self._CheckAgainstNumpy(onp_fun, jnp_fun, args_maker, check_dtypes=True)
+    self._CompileAndCheck(jnp_fun, args_maker, check_dtypes=True)
+
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_{}_axis={}".format(
           jtu.format_test_name_suffix("", [shape] * len(dtypes), dtypes), axis),