Repository: google/jax

Issue:
`random.dirichlet` sometimes samples `nan` for sparse alpha
### Description

In very particular cases, `jax.random.dirichlet` produces unexpected `nan`s. I found that these sometimes occur when some, but not all, values of `alpha` are zero. I've managed to track down one specific case. 
```python
import jax.numpy as jnp
from jax import random

# Specific key to reproduce error.
key = jnp.array([1328709358, 4232441613], dtype=jnp.uint32)
# Specific alpha, inside the zip-file attached to this issue.
# (I was unable to reproduce the error with a smaller matrix alpha, sorry!).
alpha = jnp.load('alpha.npy')
theta = random.dirichlet(key, alpha)
jnp.isnan(theta).any()  # Is True.
```
You can download `alpha.npy` from this zip: [alpha.zip](https://github.com/google/jax/files/12803309/alpha.zip)

For this specific case, `random.loggamma` in `random.dirichlet` is the culprit. Specifically, the element corresponding to `alpha[153, 38]` produces the `nan`. But only when I use the entire matrix `alpha` (not just `alpha[153]`).
```python
random.loggamma(key, alpha)[153]
```
I think this is a bug, right?
Let me know if there is anything that I can do to further clarify or help out!

Thanks in advance,

Hylke


### What jax/jaxlib version are you using?

jax-0.4.17/jaxlib-0.4.17

### Which accelerator(s) are you using?

The problem occurs both CPU & GPU.

### Additional system info

Ubuntu 23.04, Python 3.11

### NVIDIA GPU info

I don't think this is GPU specific. At any rate, the output of `nvidia-smi` on the `nvidia/cuda:12.1.1-cudnn8-devel-ubuntu22.04` Docker container:
```
+---------------------------------------------------------------------------------------+
| NVIDIA-SMI 535.54.03              Driver Version: 535.54.03    CUDA Version: 12.2     |
|-----------------------------------------+----------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |         Memory-Usage | GPU-Util  Compute M. |
|                                         |                      |               MIG M. |
|=========================================+======================+======================|
|   0  NVIDIA A40                     On  | 00000000:00:08.0 Off |                    0 |
|  0%   42C    P0              78W / 300W |  34334MiB / 46068MiB |      0%      Default |
|                                         |                      |                  N/A |
+-----------------------------------------+----------------------+----------------------+
|   1  NVIDIA A40                     On  | 00000000:00:09.0 Off |                    0 |
|  0%   41C    P0              90W / 300W |  17118MiB / 46068MiB |      0%      Default |
|                                         |                      |                  N/A |
+-----------------------------------------+----------------------+----------------------+
|   2  NVIDIA A40                     On  | 00000000:00:0A.0 Off |                    0 |
|  0%   45C    P0             115W / 300W |  17142MiB / 46068MiB |     68%      Default |
|                                         |                      |                  N/A |
+-----------------------------------------+----------------------+----------------------+
|   3  NVIDIA A40                     On  | 00000000:00:0B.0 Off |                    0 |
|  0%   41C    P0             112W / 300W |  33510MiB / 46068MiB |      0%      Default |
|                                         |                      |                  N/A |
+-----------------------------------------+----------------------+----------------------+
|   4  NVIDIA A40                     On  | 00000000:00:0C.0 Off |                    0 |
|  0%   29C    P8              21W / 300W |      4MiB / 46068MiB |      0%      Default |
|                                         |                      |                  N/A |
+-----------------------------------------+----------------------+----------------------+
|   5  NVIDIA A40                     On  | 00000000:00:0D.0 Off |                    0 |
|  0%   29C    P8              21W / 300W |      4MiB / 46068MiB |      0%      Default |
|                                         |                      |                  N/A |
+-----------------------------------------+----------------------+----------------------+
|   6  NVIDIA A40                     On  | 00000000:00:0E.0 Off |                    0 |
|  0%   31C    P8              21W / 300W |      4MiB / 46068MiB |      0%      Default |
|                                         |                      |                  N/A |
+-----------------------------------------+----------------------+----------------------+
|   7  NVIDIA A40                     On  | 00000000:00:0F.0 Off |                    0 |
|  0%   29C    P8              21W / 300W |      4MiB / 46068MiB |      0%      Default |
|                                         |                      |                  N/A |
+-----------------------------------------+----------------------+----------------------+
                                                                                         
+---------------------------------------------------------------------------------------+
| Processes:                                                                            |
|  GPU   GI   CI        PID   Type   Process name                            GPU Memory |
|        ID   ID                                                             Usage      |
|=======================================================================================|
|    0   N/A  N/A     19887      C   python                                    34322MiB |
|    1   N/A  N/A      4804      C   python                                    17106MiB |
|    2   N/A  N/A     11715      C   python                                    17130MiB |
|    3   N/A  N/A     23920      C   python                                    33498MiB |
+---------------------------------------------------------------------------------------+
```


Discussion:
Thanks for the report! It turns out this also reproduces if you use `alpha = jnp.zeros((1797, 50))`, so no need to load your specific data.

I managed to reduce the problematic case to this:
```python
import jax
from jax._src.random import _gamma_one
key=jax.random.wrap_key_data(jax.numpy.array([1057748167, 1356978999], dtype='uint32'))
print(_gamma_one(key, 0.0, log_space=True))
# nan
```
The problem is this:
```python
_, subkey = jax.random.split(key)
print(jax.random.exponential(subkey))
# 0.0
```
This happens to be a key that returns exactly zero in the exponential distribution (because it returns exactly zero in the uniform distribution), so that `log_samples` in this is zero, leading to a NaN in the following statement: https://github.com/google/jax/blob/2fe00f88a0fb336ca76cbf87e335b18e5dae1c70/jax/_src/random.py#L1171-L1172

I'll have to think about whether this is fundamentally a problem with `random.exponential`, or fundamentally a problem with the `random.loggamma` implementation.
#17925 should fix this - thanks for the report!
Wow, thanks for the incredibly fast fix! Happy I could help.

Resolving patch:
diff --git a/jax/_src/random.py b/jax/_src/random.py
--- a/jax/_src/random.py
+++ b/jax/_src/random.py
@@ -1169,7 +1169,7 @@ def _next_kxv(kxv):
   _, _, V, _ = lax.while_loop(_cond_fn, _body_fn, (key, zero, one, _lax_const(alpha, 2)))
   if log_space:
     log_samples = lax.neg(exponential(subkey, (), dtype=dtype))
-    log_boost = lax.select(boost_mask, zero, lax.mul(log_samples, lax.div(one, alpha_orig)))
+    log_boost = lax.select(boost_mask | (log_samples == 0), zero, lax.mul(log_samples, lax.div(one, alpha_orig)))
     return lax.add(lax.add(lax.log(d), lax.log(V)), log_boost)
   else:
     samples = 1 - uniform(subkey, (), dtype=dtype)


Test patch:
diff --git a/tests/random_test.py b/tests/random_test.py
--- a/tests/random_test.py
+++ b/tests/random_test.py
@@ -451,6 +451,20 @@ def test_threefry_split_vmapped_fold_in_symmetry(self, make_key):
       self.assertArraysEqual(f2, s2)
       self.assertArraysEqual(f3, s3)
 
+  @skipIf(config.jax_threefry_partitionable, 'changed random bit values')
+  def test_loggamma_nan_corner_case(self):
+    # regression test for https://github.com/google/jax/issues/17922
+    # This particular key previously led to NaN output.
+    # If the underlying implementation ever changes, this test will no longer
+    # exercise this corner case, so we compare to a particular output value
+    # rather than just checking for lack of NaNs.
+    expected = jnp.float32(-4.595436)
+    key = random.wrap_key_data(
+      jnp.array([3200590325, 713258242], dtype='uint32'))
+    actual = random.loggamma(key, 0.0, dtype='float32')
+    rtol = 1E-4 if jtu.test_device_matches(["tpu"]) else 1E-6
+    self.assertAllClose(expected, actual, rtol=rtol)
+
   @parameterized.parameters([params
       for d in [
           {"seed": 0, "typ": int, "jit": True, "key": [0, 0]},