Repository: google/jax

Issue:
jax.numpy.median throws error for float median
`import jax.numpy as np; np.median(np.array([i for i in range(0,10)]))` throws an `TypeError: q and a arguments to quantile must be of float type, got int32 and float3` where `import numpy as np; np.median(np.array([i for i in range(0,10)]))` results in `4.5`.

<details><summary>Traceback of error (Click!)</summary>
<p>

```python
TypeError                                 Traceback (most recent call last)
<ipython-input-28-eb17bd577cca> in <module>
----> 1 np.median(np.array([i for i in range(0,10)]))

/opt/anaconda3/lib/python3.7/site-packages/jax/numpy/lax_numpy.py in median(a, axis, out, overwrite_input, keepdims)
   3556     q = 0.5
   3557     return quantile(a, q, axis=axis, out=out, overwrite_input=overwrite_input,
-> 3558                     keepdims=keepdims, interpolation='midpoint')
   3559 
   3560 def _astype(arr, dtype):

/opt/anaconda3/lib/python3.7/site-packages/jax/numpy/lax_numpy.py in quantile(a, q, axis, out, overwrite_input, interpolation, keepdims)
   3464   if interpolation not in ["linear", "lower", "higher", "midpoint", "nearest"]:
   3465     raise ValueError("interpolation can only be 'linear', 'lower', 'higher', 'midpoint', or 'nearest'")
-> 3466   return _quantile(a, q, axis, interpolation, keepdims)
   3467 
   3468 @partial(jit, static_argnums=(2, 3, 4))

/opt/anaconda3/lib/python3.7/site-packages/jax/api.py in f_jitted(*args, **kwargs)
    151     flat_fun, out_tree = flatten_fun(f, in_tree)
    152     out = xla.xla_call(flat_fun, *args_flat, device=device, backend=backend,
--> 153                        name=flat_fun.__name__)
    154     return tree_unflatten(out_tree(), out)
    155 

/opt/anaconda3/lib/python3.7/site-packages/jax/core.py in call_bind(primitive, f, *args, **params)
    976   if top_trace is None:
    977     with new_sublevel():
--> 978       outs = primitive.impl(f, *args, **params)
    979   else:
    980     tracers = map(top_trace.full_raise, args)

/opt/anaconda3/lib/python3.7/site-packages/jax/interpreters/xla.py in _xla_call_impl(fun, device, backend, name, *args)
    461 
    462 def _xla_call_impl(fun: lu.WrappedFun, *args, device, backend, name):
--> 463   compiled_fun = _xla_callable(fun, device, backend, name, *map(arg_spec, args))
    464   try:
    465     return compiled_fun(*args)

/opt/anaconda3/lib/python3.7/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

/opt/anaconda3/lib/python3.7/site-packages/jax/interpreters/xla.py in _xla_callable(fun, device, backend, name, *arg_specs)
    478   pvals: Sequence[pe.PartialVal] = [pe.PartialVal.unknown(aval) for aval in abstract_args]
    479   jaxpr, pvals, consts = pe.trace_to_jaxpr(
--> 480       fun, pvals, instantiate=False, stage_out=True, bottom=True)
    481 
    482   _map(prefetch, it.chain(consts, jaxpr_literals(jaxpr)))

/opt/anaconda3/lib/python3.7/site-packages/jax/interpreters/partial_eval.py in trace_to_jaxpr(fun, pvals, instantiate, stage_out, bottom)
    419   with new_master(trace_type, bottom=bottom) as master:
    420     fun = trace_to_subjaxpr(fun, master, instantiate)
--> 421     jaxpr, (out_pvals, consts, env) = fun.call_wrapped(pvals)
    422     assert not env
    423     del master

/opt/anaconda3/lib/python3.7/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:

/opt/anaconda3/lib/python3.7/site-packages/jax/numpy/lax_numpy.py in _quantile(a, q, axis, interpolation, keepdims)
   3485   if not issubdtype(a.dtype, floating) or not issubdtype(q.dtype, floating):
   3486     msg = "q and a arguments to quantile must be of float type, got {} and {}"
-> 3487     raise TypeError(msg.format(a.dtype, q.dtype))
   3488 
   3489   # Promote q to at least float32 for precise interpolation.

TypeError: q and a arguments to quantile must be of float type, got int32 and float32
```

</p>
</details>

jax 0.1.63, Python 3.7.4


Discussion:
Hello @Baschdl 
Solution to this, 
```
from jax import numpy as np
np.median(np.array([i for i in range(0,10)],dtype = np.float32))
```

You just need to add **dtype** in it.

Regards.
Shorter repro:
```
import jax.numpy as jnp
jnp.median(jnp.arange(10))
```
I'll work on a fix.

Resolving patch:
diff --git a/jax/numpy/lax_numpy.py b/jax/numpy/lax_numpy.py
--- a/jax/numpy/lax_numpy.py
+++ b/jax/numpy/lax_numpy.py
@@ -3790,7 +3790,8 @@ def quantile(a, q, axis=None, out=None, overwrite_input=False,
 
 @partial(jit, static_argnums=(2, 3, 4))
 def _quantile(a, q, axis, interpolation, keepdims):
-  a = asarray(a)
+  a = asarray(a, dtype=promote_types(_dtype(a), float_))
+  q = asarray(q, dtype=promote_types(_dtype(q), float_))
   if axis is None:
     a = ravel(a)
     axis = 0
@@ -3803,15 +3804,6 @@ def _quantile(a, q, axis, interpolation, keepdims):
   if q_ndim > 1:
     raise ValueError("q must be have rank <= 1, got shape {}".format(shape(q)))
 
-  q = asarray(q)
-
-  if not issubdtype(a.dtype, floating) or not issubdtype(q.dtype, floating):
-    msg = "q and a arguments to quantile must be of float type, got {} and {}"
-    raise TypeError(msg.format(a.dtype, q.dtype))
-
-  # Promote q to at least float32 for precise interpolation.
-  q = lax.convert_element_type(q, promote_types(q.dtype, float32))
-
   a_shape = shape(a)
   a = lax.sort(a, dimension=axis)
 


Test patch:
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -2434,7 +2434,7 @@ def args_maker(): return []
           ("percentile", partial(jtu.rand_uniform, low=0., high=100.)),
           ("quantile", partial(jtu.rand_uniform, low=0., high=1.)),
         )
-        for a_dtype in float_dtypes
+        for a_dtype in default_dtypes
         for a_shape, axis in (
           ((7,), None),
           ((47, 7), 0),
@@ -2477,7 +2477,7 @@ def onp_fun(*args):
          "a_shape": a_shape, "a_dtype": a_dtype,
          "axis": axis,
          "keepdims": keepdims}
-        for a_dtype in float_dtypes
+        for a_dtype in default_dtypes
         for a_shape, axis in (
           ((7,), None),
           ((47, 7), 0),