Repository: google/jax

Issue:
Cannot specify data type to be int in `random.normal`
I have been trying to use `random` with data type `np.int32` or any other int type (e.g `np.int8`, `np.in16`). 

Recreating the bug:

```python
from jax import random
import jax.numpy as np

key = random.PRNGKey(4)
x = random.normal(key,dtype=np.int32)
``` 

Error message:
```
TypeError: No loop matching the specified signature and casting was found for ufunc nextafter
```

Full logs can be found below:
```
TypeError                                 Traceback (most recent call last)
<ipython-input-1-eb3ab6450cb1> in <module>
      3 
      4 key = random.PRNGKey(4)
----> 5 x = random.normal(key,dtype=np.int32)

~/anaconda3/lib/python3.6/site-packages/jax/random.py in normal(key, shape, dtype)
    546   dtype = dtypes.canonicalize_dtype(dtype)
    547   shape = abstract_arrays.canonicalize_shape(shape)
--> 548   return _normal(key, shape, dtype)
    549 
    550 @partial(jit, static_argnums=(1, 2))

~/anaconda3/lib/python3.6/site-packages/jax/api.py in f_jitted(*args, **kwargs)
    165     flat_fun, out_tree = flatten_fun(f, in_tree)
    166     out = xla.xla_call(flat_fun, *args_flat, device=device, backend=backend,
--> 167                        name=flat_fun.__name__, donated_invars=donated_invars)
    168     return tree_unflatten(out_tree(), out)
    169 

~/anaconda3/lib/python3.6/site-packages/jax/core.py in _call_bind(processor, post_processor, primitive, f, *args, **params)
   1077   if top_trace is None:
   1078     with new_sublevel():
-> 1079       outs = primitive.impl(f, *args, **params)
   1080   else:
   1081     tracers = map(top_trace.full_raise, args)

~/anaconda3/lib/python3.6/site-packages/jax/interpreters/xla.py in _xla_call_impl(fun, device, backend, name, donated_invars, *args)
    525 
    526 def _xla_call_impl(fun: lu.WrappedFun, *args, device, backend, name, donated_invars):
--> 527   compiled_fun = _xla_callable(fun, device, backend, name, donated_invars, *map(arg_spec, args))
    528   try:
    529     return compiled_fun(*args)

~/anaconda3/lib/python3.6/site-packages/jax/linear_util.py in memoized_fun(fun, *args)
    219       fun.populate_stores(stores)
    220     else:
--> 221       ans = call(fun, *args)
    222       cache[key] = (ans, fun.stores)
    223     return ans

~/anaconda3/lib/python3.6/site-packages/jax/interpreters/xla.py in _xla_callable(fun, device, backend, name, donated_invars, *arg_specs)
    585   pvals: Sequence[pe.PartialVal] = [pe.PartialVal.unknown(aval) for aval in abstract_args]
    586   jaxpr, pvals, consts = pe.trace_to_jaxpr(
--> 587       fun, pvals, instantiate=False, stage_out=True, bottom=True)
    588   jaxpr, uses_outfeed = apply_outfeed_rewriter(jaxpr)
    589   _map(prefetch, it.chain(consts, jaxpr_literals(jaxpr)))

~/anaconda3/lib/python3.6/site-packages/jax/interpreters/partial_eval.py in trace_to_jaxpr(fun, pvals, instantiate, stage_out, bottom, trace_type)
    450   with new_master(trace_type, bottom=bottom) as master:
    451     fun = trace_to_subjaxpr(fun, master, instantiate)
--> 452     jaxpr, (out_pvals, consts, env) = fun.call_wrapped(pvals)
    453     assert not env
    454     del master

~/anaconda3/lib/python3.6/site-packages/jax/linear_util.py in call_wrapped(self, *args, **kwargs)
    148     gen = None
    149 
--> 150     ans = self.f(*args, **dict(self.params, **kwargs))
    151     del args
    152     while stack:

~/anaconda3/lib/python3.6/site-packages/jax/random.py in _normal(key, shape, dtype)
    551 def _normal(key, shape, dtype):
    552   _check_shape("normal", shape)
--> 553   lo = np.nextafter(np.array(-1., dtype), 0., dtype=dtype)
    554   hi = np.array(1., dtype)
    555   u = uniform(key, shape, dtype, lo, hi)

TypeError: No loop matching the specified signature and casting was found for ufunc nextafter
```


Discussion:
Thanks for reporting this!

I see this as an issue of raising a bad error message; we probably want the `random.normal` API only to accept inexect (i.e. floating point) dtypes and not int dtypes, but one can always cast the output to an integer as needed.
Thanks a lot for your prompt comment!

Yes, that's correct. I have also ended up casting these values to ints. However, just noticed I cannot use `dtype=np.complex128` or `dtype=np.complex64`. Is it still part of intended behavior of this API?
I think this is probably intended; by analogy, `np.random.normal` doesn't accept dtype at all, and can only return float64. If you want to construct complex normal values you need to do it by way of floats.
> I think this is probably intended; by analogy, `np.random.normal` doesn't accept dtype at all, and can only return float64. If you want to construct complex normal values you need to do it by way of floats.

That's very true. I realised my previous response was irrelevant, hence deleted it. 

Also, I figured one can specify dtype using `astype`, instead of third parameter in `jax.random.normal`. 

```python
from jax import random
import jax.numpy as np

key = random.PRNGKey(4)
x = random.normal(key,(1,)).astype(np.int32)
```
This only works if I am generating an array, but not a single value. 

Resolving patch:
diff --git a/jax/random.py b/jax/random.py
--- a/jax/random.py
+++ b/jax/random.py
@@ -359,6 +359,9 @@ def uniform(key: jnp.ndarray,
   Returns:
     A random array with the specified shape and dtype.
   """
+  if not dtypes.issubdtype(dtype, np.floating):
+    raise ValueError(f"dtype argument to `uniform` must be a float dtype, "
+                     f"got {dtype}")
   dtype = dtypes.canonicalize_dtype(dtype)
   shape = abstract_arrays.canonicalize_shape(shape)
   return _uniform(key, shape, dtype, minval, maxval)
@@ -543,6 +546,9 @@ def normal(key: jnp.ndarray,
   Returns:
     A random array with the specified shape and dtype.
   """
+  if not dtypes.issubdtype(dtype, np.floating):
+    raise ValueError(f"dtype argument to `normal` must be a float dtype, "
+                     f"got {dtype}")
   dtype = dtypes.canonicalize_dtype(dtype)
   shape = abstract_arrays.canonicalize_shape(shape)
   return _normal(key, shape, dtype)
@@ -581,6 +587,9 @@ def multivariate_normal(key: jnp.ndarray,
     ``shape + mean.shape[-1:]`` if ``shape`` is not None, or else
     ``broadcast_shapes(mean.shape[:-1], cov.shape[:-2]) + mean.shape[-1:]``.
   """
+  if not dtypes.issubdtype(dtype, np.floating):
+    raise ValueError(f"dtype argument to `multivariate_normal` must be a float "
+                     f"dtype, got {dtype}")
   dtype = dtypes.canonicalize_dtype(dtype)
   if shape is not None:
     shape = abstract_arrays.canonicalize_shape(shape)
@@ -634,6 +643,9 @@ def truncated_normal(key: jnp.ndarray,
     A random array with the specified dtype and shape given by ``shape`` if
     ``shape`` is not None, or else by broadcasting ``lower`` and ``upper``.
   """
+  if not dtypes.issubdtype(dtype, np.floating):
+    raise ValueError(f"dtype argument to `truncated_normal` must be a float "
+                     f"dtype, got {dtype}")
   dtype = dtypes.canonicalize_dtype(dtype)
   if shape is not None:
     shape = abstract_arrays.canonicalize_shape(shape)
@@ -714,6 +726,9 @@ def beta(key: jnp.ndarray,
     A random array with the specified dtype and shape given by ``shape`` if
     ``shape`` is not None, or else by broadcasting ``a`` and ``b``.
   """
+  if not dtypes.issubdtype(dtype, np.floating):
+    raise ValueError(f"dtype argument to `beta` must be a float "
+                     f"dtype, got {dtype}")
   dtype = dtypes.canonicalize_dtype(dtype)
   if shape is not None:
     shape = abstract_arrays.canonicalize_shape(shape)
@@ -748,6 +763,9 @@ def cauchy(key, shape=(), dtype=np.float64):
   Returns:
     A random array with the specified shape and dtype.
   """
+  if not dtypes.issubdtype(dtype, np.floating):
+    raise ValueError(f"dtype argument to `cauchy` must be a float "
+                     f"dtype, got {dtype}")
   dtype = dtypes.canonicalize_dtype(dtype)
   shape = abstract_arrays.canonicalize_shape(shape)
   return _cauchy(key, shape, dtype)
@@ -780,6 +798,9 @@ def dirichlet(key, alpha, shape=None, dtype=np.float64):
     ``shape + (alpha.shape[-1],)`` if ``shape`` is not None, or else
     ``alpha.shape``.
   """
+  if not dtypes.issubdtype(dtype, np.floating):
+    raise ValueError(f"dtype argument to `dirichlet` must be a float "
+                     f"dtype, got {dtype}")
   dtype = dtypes.canonicalize_dtype(dtype)
   if shape is not None:
     shape = abstract_arrays.canonicalize_shape(shape)
@@ -814,6 +835,9 @@ def exponential(key, shape=(), dtype=np.float64):
   Returns:
     A random array with the specified shape and dtype.
   """
+  if not dtypes.issubdtype(dtype, np.floating):
+    raise ValueError(f"dtype argument to `exponential` must be a float "
+                     f"dtype, got {dtype}")
   dtype = dtypes.canonicalize_dtype(dtype)
   shape = abstract_arrays.canonicalize_shape(shape)
   return _exponential(key, shape, dtype)
@@ -1039,6 +1063,9 @@ def gamma(key, a, shape=None, dtype=np.float64):
     A random array with the specified dtype and with shape given by ``shape`` if
     ``shape`` is not None, or else by ``a.shape``.
   """
+  if not dtypes.issubdtype(dtype, np.floating):
+    raise ValueError(f"dtype argument to `gamma` must be a float "
+                     f"dtype, got {dtype}")
   dtype = dtypes.canonicalize_dtype(dtype)
   if shape is not None:
     shape = abstract_arrays.canonicalize_shape(shape)
@@ -1178,6 +1205,9 @@ def gumbel(key, shape=(), dtype=np.float64):
   Returns:
     A random array with the specified shape and dtype.
   """
+  if not dtypes.issubdtype(dtype, np.floating):
+    raise ValueError(f"dtype argument to `gumbel` must be a float "
+                     f"dtype, got {dtype}")
   dtype = dtypes.canonicalize_dtype(dtype)
   shape = abstract_arrays.canonicalize_shape(shape)
   return _gumbel(key, shape, dtype)
@@ -1232,6 +1262,9 @@ def laplace(key, shape=(), dtype=np.float64):
   Returns:
     A random array with the specified shape and dtype.
   """
+  if not dtypes.issubdtype(dtype, np.floating):
+    raise ValueError(f"dtype argument to `laplace` must be a float "
+                     f"dtype, got {dtype}")
   dtype = dtypes.canonicalize_dtype(dtype)
   shape = abstract_arrays.canonicalize_shape(shape)
   return _laplace(key, shape, dtype)
@@ -1257,6 +1290,9 @@ def logistic(key, shape=(), dtype=np.float64):
   Returns:
     A random array with the specified shape and dtype.
   """
+  if not dtypes.issubdtype(dtype, np.floating):
+    raise ValueError(f"dtype argument to `logistic` must be a float "
+                     f"dtype, got {dtype}")
   dtype = dtypes.canonicalize_dtype(dtype)
   shape = abstract_arrays.canonicalize_shape(shape)
   return _logistic(key, shape, dtype)
@@ -1297,6 +1333,9 @@ def pareto(key, b, shape=None, dtype=np.float64):
     A random array with the specified dtype and with shape given by ``shape`` if
     ``shape`` is not None, or else by ``b.shape``.
   """
+  if not dtypes.issubdtype(dtype, np.floating):
+    raise ValueError(f"dtype argument to `pareto` must be a float "
+                     f"dtype, got {dtype}")
   dtype = dtypes.canonicalize_dtype(dtype)
   if shape is not None:
     shape = abstract_arrays.canonicalize_shape(shape)
@@ -1331,6 +1370,9 @@ def t(key, df, shape=(), dtype=np.float64):
     A random array with the specified dtype and with shape given by ``shape`` if
     ``shape`` is not None, or else by ``df.shape``.
   """
+  if not dtypes.issubdtype(dtype, np.floating):
+    raise ValueError(f"dtype argument to `t` must be a float "
+                     f"dtype, got {dtype}")
   dtype = dtypes.canonicalize_dtype(dtype)
   shape = abstract_arrays.canonicalize_shape(shape)
   return _t(key, df, shape, 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
@@ -687,6 +687,10 @@ def testPRNGValues(self):
         random.fold_in(k, 4),
         np.array([2285895361,  433833334], dtype='uint32'))
 
+  def testDtypeErrorMessage(self):
+    with self.assertRaisesRegex(ValueError, r"dtype argument to.*"):
+      random.normal(random.PRNGKey(0), (), dtype=jnp.int32)
+
 
 if __name__ == "__main__":
   absltest.main()