Merge pull request #9564 from MichaelMarien:random-choice-docstring

PiperOrigin-RevId: 428596963
jax-ml · Feb 14, 2022 · d569440 · d569440
2 parents beaa00c + 20e5090
commit d569440
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/jax/_src/random.py b/jax/_src/random.py
@@ -435,6 +435,11 @@ def choice(key: KeyArray,
            axis: int = 0) -> jnp.ndarray:
   """Generates a random sample from a given array.
 
+  .. warning::
+    If ``p`` has fewer non-zero elements than the requested number of samples,
+    as specified in ``shape``, and ``replace=False``, the output of this
+    function is ill-defined. Please make sure to use appropriate inputs.
+
   Args:
     key: a PRNG key used as the random key.
     a : array or int. If an ndarray, a random sample is generated from