Inconsistent dtype output in jax.random.choice for double precision enabled
### Description

When using double precision, `jax.random.choice` returns inconsistent dtypes. When using a weight vector, the output dtype is set using `searchsorted` rules, i.e., based on the size of the array of weights at hand, while when none is used (uniform weights), the dtype is set using `jax.random.randint` rules (defaulting to `jnp.int_`). I believe the behaviour should be homogenous, no matter what the choice ends up being.

```python
import jax
import jax.numpy as jnp

jax.config.update("jax_enable_x64", True)
key = jax.random.PRNGKey(42)
w = jnp.array([1/3, 1/3, 1/3])

print(jax.random.choice(key, 3).dtype)
print(jax.random.choice(key, 3, p=w).dtype)
```

In practice, I believe that the implementation choice made in https://github.com/google/jax/blob/c22302ae2e0913efd5be23e136f2ff14b520d9cf/jax/_src/numpy/lax_numpy.py#L4824 is perhaps a bit redundant. Indeed, if the config has not enables double precision, `int64` will default to `int32` anyway, while, if it has, you were happy to pay the price in the first place.

WDYT?



### What jax/jaxlib version are you using?

jax v0.4.10 

### Which accelerator(s) are you using?

CPU

### Additional system info

_No response_

### NVIDIA GPU info

_No response_
diff --git a/jax/_src/random.py b/jax/_src/random.py
index 08d64dadc..56cc53cdf 100644
--- a/jax/_src/random.py
+++ b/jax/_src/random.py
@@ -567,7 +567,7 @@ def choice(key: KeyArray,
     if replace:
       p_cuml = jnp.cumsum(p_arr)
       r = p_cuml[-1] * (1 - uniform(key, shape, dtype=p_cuml.dtype))
-      ind = jnp.searchsorted(p_cuml, r)
+      ind = jnp.searchsorted(p_cuml, r).astype(int)
     else:
       # Gumbel top-k trick: https://timvieira.github.io/blog/post/2019/09/16/algorithms-for-sampling-without-replacement/
       g = -gumbel(key, (n_inputs,), dtype=p_arr.dtype) - jnp.log(p_arr)
diff --git a/tests/random_test.py b/tests/random_test.py
index 26a544688..93fc96c01 100644
--- a/tests/random_test.py
+++ b/tests/random_test.py
@@ -754,8 +754,10 @@ class LaxRandomTest(jtu.JaxTestCase):
     sample = rand(key, x)
     if not is_range:
       self.assertEqual(dtype, sample.dtype)
-    np_shape = np.shape(np_choice(x, shape or None, replace, p, axis))
-    self.assertEqual(np_shape, sample.shape)
+    expected_shape = np.shape(np_choice(x, shape or None, replace, p, axis))
+    self.assertEqual(expected_shape, sample.shape)
+    expected_dtype = dtypes.result_type(int if is_range else x)
+    self.assertEqual(expected_dtype, sample.dtype)
     if not replace and shape:
       def lsort(x):
         if not math.prod(x.shape): return x
['jax/_src/random.py', 'tests/random_test.py']