Repository: jax-ml/jax (python)

Issue:
bernoulli samples fail for small probabilities
Example:
```python
import jax

key = jax.random.key(0)
samples = jax.random.bernoulli(key, p=1E-10, shape=int(1E8))
print(samples.sum())  # should be zero with high probability
# 7
```
The issue is that at float32, uniform values have a resolution of $2^{-23}$ which is about $10^{-7}$, and so we don't have the resolution to properly evaluate samples with `p` smaller than about `1E-7`.

We could make this better by implementing bernoulli sampling directly on uint32 bits rather than converting to float, though that would only get us to `p` of about `2E-10`. By implementing an optional super-sampling (i.e. splitting the key) we could get to `4E-10`.

Pull request description:
jax.random.bernoulli: use higher-resolution sampler
Fixes #28017

This approach works sufficiently for `p` down to about `1E-10`, compared to the current approach which is good for `p` down to about `1E-7`.

To support smaller `p` than this, we'll need to add a `mode` argument similar to what we have in [`jax.random.gumbel`](https://docs.jax.dev/en/latest/_autosummary/jax.random.gumbel.html).

In a happy accident of how uniform samples are generated, this new approach produces the same samples for a given key as the old approach, except when `p` is very small (which is the part we wanted to change).

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index d46d9d01a..3bd701777 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -43,6 +43,9 @@ When releasing, please add the new-release-boilerplate to docs/pallas/CHANGELOG.
   * {func}`jax.jit` now requires `fun` to be passed by position, and additional
     arguments to be passed by keyword. Doing otherwise will result in a
     DeprecationWarning in v0.6.X, and an error in starting in v0.7.X.
+  * {func}`jax.random.beroulli` now has higher resolution, and can correctly handle
+    values of `p` down to about `1E-10`. Previously results were incorrect for `p`
+    smaller than about `1E-7`. ({jax-issue}`#28022`)
 
 * Deprecations
 
diff --git a/jax/_src/random.py b/jax/_src/random.py
index b29c1dca7..b50a9653a 100644
--- a/jax/_src/random.py
+++ b/jax/_src/random.py
@@ -958,7 +958,13 @@ def _bernoulli(key, p, shape) -> Array:
   else:
     _check_shape("bernoulli", shape, np.shape(p))
 
-  return uniform(key, shape, lax.dtype(p)) < p
+  # we could return uniform(key, shape) < p. But uniform sacrifices some
+  # resolution, so instead we sample in the space of the generated bits.
+  nbits = dtypes.finfo(p.dtype).bits
+  unsigned_dtype = UINT_DTYPES[nbits]
+  samples = bits(key, shape, unsigned_dtype)
+  cutoff = (p * _lax_const(p, 1 << nbits)).astype(unsigned_dtype)
+  return (p > 0) & (samples < cutoff) | (p >= 1)
 
 
 def beta(key: ArrayLike,


Test patch:
diff --git a/tests/random_lax_test.py b/tests/random_lax_test.py
index b4d2853ab..6230b6471 100644
--- a/tests/random_lax_test.py
+++ b/tests/random_lax_test.py
@@ -461,6 +461,12 @@ class DistributionsTest(RandomTestBase):
       x = random.bernoulli(key, np.array([0.2, 0.3]), shape=(3, 2))
     assert x.shape == (3, 2)
 
+  def testBernoulliSmallProbabilty(self):
+    # Regression test for https://github.com/jax-ml/jax/issues/28017
+    key = jax.random.key(0)
+    samples = jax.random.bernoulli(key, p=1E-10, shape=int(1E8))
+    self.assertEqual(samples.sum(), 0)
+
   @jtu.sample_product(
     a=[0.2, 5.],
     b=[0.2, 5.],