Repository: jax-ml/jax
Pull request: Update random.logistic() to prevent infinities
Fixes #3043

Problem statement:
random.logistic returns -inf
If I call jax.random.logistic with key array([2308802225, 3876068305], dtype=uint32) and shape=(10000, 12), the sampled array contains -inf. I don't think this behavior is correct or at least it is not expected and causes downstream code that assumes finite floats to crash.

I can work around this, but I think it's worth fixing. Thanks!

Hints from issue discussion:
This is a consequence of using a 32-bit implementation of the standard logit generating function: x is drawn from a uniform distribution between 0 and 1, then transformed according to `y = log(x / (1 - x))`. For real numbers there is an infinitessimally small chance of `x=0` which maps to `y=-inf`, in 32-bit arithmetic the chance is approximately 1 in 2^30.

I'll think about how to best address this.
Wait, it's even worse than that, because ``uniform()`` only randomizes the mantissa: this means the chance is about a 1 in 2^23 of getting -inf; that's confirmed empirically:
```
>>> (uniform(PRNGKey(0), (100000000,)) == 0).sum()
DeviceArray(19, dtype=int32)
```
I think the best approach would be to modify the generating algorithm to use 32-bit unsigned integers rather than 32-bit uniform floats; then the chances of hitting this would be one in 4 billion, rather than 1 in 8 million as it is currently.
Could you avoid the problem by drawing from the range `(0, 1)` instead of `[0, 1)` ?
Yes, that would fix the -inf issue. But it still might be preferable to use an approach that allows the result to have a cardinality greater than 2^23

Commit message: Make comment more concise

Patch to review:
diff --git a/jax/random.py b/jax/random.py
index 13eae646c46d..0227f309c17e 100644
--- a/jax/random.py
+++ b/jax/random.py
@@ -1227,8 +1227,21 @@ def logistic(key, shape=(), dtype=onp.float64):
 
 @partial(jit, static_argnums=(1, 2))
 def _logistic(key, shape, dtype):
+  # Mathematically, we can compute the distribution by generating uniformly-distributed
+  # numbers x in the open interval (a, b) and computing:
+  #   z = log[ (x - a) / (b - x))
+  # It's important to avoid x=a or x=b, which lead to infinite values for z.
+  # The uniform() function genreates pseudorandom floating point numbers x in the
+  # semi-closed interval [0, 1), so if used directly  with (a,b)=(0,1), it will
+  # lead to infinite output in a small number of cases (as many as 1 in 2^23 for float32).
+  #
+  # Instead, we let (a, b) = (-ε, 1) where ε is the smallest step between floating point
+  # values: then numbers in the interval (-ε, 1) are approximated by standard uniformly
+  # drawn numbers in [0, 1).
   _check_shape("logistic", shape)
-  return logit(uniform(key, shape, dtype))
+  x = uniform(key, shape, dtype)
+  eps = np.finfo(dtype).eps
+  return lax.log(lax.div(lax.add(lax._const(x, eps), x), lax.sub(lax._const(x, 1), x)))
 
 
 def pareto(key, b, shape=None, dtype=onp.float64):


Code review comments:
diff_hunk: @@ -1227,8 +1227,21 @@ def logistic(key, shape=(), dtype=onp.float64):
 
 @partial(jit, static_argnums=(1, 2))
 def _logistic(key, shape, dtype):
+  # Mathematically, we can compute the distribution by generating uniformly-distributed
+  # numbers x in the open interval (a, b) and computing:
+  #   z = log[ (x - a) / (b - x))
+  # It's important to avoid x=a or x=b, which lead to infinite values for z.
+  # The uniform() function genreates pseudorandom floating point numbers x in the
path: jax/random.py
text: @user1:
sp. generates

@author:
Done.