Repository: google/jax

Issue:
Incorrect type promotion in random.choice
As discovered in #11237, `random.choice` throws a type promotion error when x64 is enabled and `p` is `float32`.

The following demonstrates the issue. Executing:

```python
import jax
import jax.numpy as jnp
jax.config.update("jax_enable_x64", True)
with jax.numpy_dtype_promotion("strict"):
    jax.random.choice(jax.random.PRNGKey(0), 5, p=jnp.ones(5, dtype=jnp.float32))
```

results in the following type promotion error:

<details>
<summary>Traceback</summary>

```python
Traceback (most recent call last):
  File "<stdin>", line 2, in <module>
  File "/Users/dforemanmackey/src/google/jax/jax/_src/random.py", line 499, in choice
    r = p_cuml[-1] * (1 - uniform(key, shape))
  File "/Users/dforemanmackey/src/google/jax/jax/_src/numpy/lax_numpy.py", line 4586, in deferring_binary_op
    return binary_op(self, other)
  File "/Users/dforemanmackey/src/google/jax/jax/_src/numpy/ufuncs.py", line 80, in fn
    x1, x2 = _promote_args(numpy_fn.__name__, x1, x2)
  File "/Users/dforemanmackey/src/google/jax/jax/_src/numpy/util.py", line 332, in _promote_args
    return _promote_shapes(fun_name, *_promote_dtypes(*args))
  File "/Users/dforemanmackey/src/google/jax/jax/_src/numpy/util.py", line 262, in _promote_dtypes
    to_dtype, weak_type = dtypes._lattice_result_type(*args)
jax._src.dtypes.TypePromotionError: Input dtypes ('float32', 'float64') have no available implicit dtype promotion path when jax_numpy_dtype_promotion=strict. Try explicitly casting inputs to the desired output type, or set jax_numpy_dtype_promotion=standard.
```
</details>

The offending line is: https://github.com/google/jax/blob/de464fcf22c8bf7a2931182f8095bc01530df9fe/jax/_src/random.py#L499

And I think the fix is as simple as:

```diff
- r = p_cuml[-1] * (1 - uniform(key, shape))
+ r = p_cuml[-1] * (1 - uniform(key, shape, dtype=p_cuml.dtype))
```

But I haven't checked in detail!


Discussion:
Good catch - thanks!

Resolving patch:
diff --git a/jax/_src/random.py b/jax/_src/random.py
--- a/jax/_src/random.py
+++ b/jax/_src/random.py
@@ -30,7 +30,7 @@
 from jax._src.api import jit, vmap
 from jax._src.lax import lax as lax_internal
 from jax._src.lib import xla_bridge
-from jax._src.numpy.lax_numpy import _arraylike, _check_arraylike, _convert_and_clip_integer
+from jax._src.numpy.lax_numpy import _arraylike, _check_arraylike, _convert_and_clip_integer, _promote_dtypes_inexact
 from jax.numpy.linalg import cholesky, svd, eigh
 from jax.interpreters import ad
 from jax.interpreters import batching
@@ -492,15 +492,17 @@ def choice(key: KeyArray,
       slices = (slice(None),) * axis + (slice(n_draws),)
       result = permutation(key, a, axis)[slices]
   else:
+    _check_arraylike("choice", p)
+    p, = _promote_dtypes_inexact(p)
     if p.shape != (n_inputs,):
       raise ValueError("p must be None or match the shape of a")
     if replace:
       p_cuml = jnp.cumsum(p)
-      r = p_cuml[-1] * (1 - uniform(key, shape))
+      r = p_cuml[-1] * (1 - uniform(key, shape, dtype=p_cuml.dtype))
       ind = jnp.searchsorted(p_cuml, r)
     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)
+      g = -gumbel(key, (n_inputs,), dtype=p.dtype) - jnp.log(p)
       ind = jnp.argsort(g)[:n_draws]
     result = ind if np.ndim(a) == 0 else jnp.take(a, ind, axis)
 


Test patch:
diff --git a/tests/random_test.py b/tests/random_test.py
--- a/tests/random_test.py
+++ b/tests/random_test.py
@@ -28,12 +28,12 @@
 
 import jax
 from jax import core
-from jax import dtypes
 from jax import grad
 from jax import lax
 from jax import numpy as jnp
 from jax import prng
 from jax import random
+from jax._src import dtypes
 from jax._src import test_util as jtu
 from jax import vmap
 from jax.interpreters import xla
@@ -639,6 +639,7 @@ def testShuffle(self, dtype):
   def testChoice(self, dtype, input_range_or_shape, shape, replace, weighted, axis):
     # This is the function API that we test against (note that self.rng().choice differs)
     np_choice = np.random.default_rng(0).choice
+    p_dtype = dtypes._to_inexact_dtype(dtype)
 
     key = self.seed_prng(0)
     is_range = type(input_range_or_shape) is int
@@ -646,7 +647,11 @@ def testChoice(self, dtype, input_range_or_shape, shape, replace, weighted, axis
          self.rng().permutation(np.arange(np.prod(
            input_range_or_shape), dtype=dtype)).reshape(input_range_or_shape))
     N = x if is_range else x.shape[axis]
-    p = None if not weighted else (np.arange(N) + 1) / np.sum(np.arange(N) + 1)
+    if weighted:
+      p = np.arange(N, dtype=p_dtype) + 1
+      p /= p.sum()
+    else:
+      p = None
     rand = lambda key, x: random.choice(key, x, shape, replace, p, axis)
     sample = rand(key, x)
     if not is_range: