# Buggy code
def _quantile(a, q, axis, interpolation, keepdims):
    a = asarray(a)
    if axis is None:
        a = ravel(a)
        axis = 0
    elif isinstance(axis, tuple):
        raise NotImplementedError("Tuple values for axis are not implemented")
    else:
        axis = _canonicalize_axis(axis, ndim(a))

    q_ndim = ndim(q)
    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)

    n = a_shape[axis]
    q = lax.mul(q, _constant_like(q, n - 1))
    low = lax.floor(q)
    high = lax.ceil(q)
    high_weight = lax.sub(q, low)
    low_weight = lax.sub(_constant_like(high_weight, 1), high_weight)

    low = lax.clamp(_constant_like(low, 0), low, _constant_like(low, n - 1))
    high = lax.clamp(_constant_like(high, 0), high, _constant_like(high, n - 1))
    low = lax.convert_element_type(low, int64)
    high = lax.convert_element_type(high, int64)

    slice_sizes = list(a_shape)
    slice_sizes[axis] = 1

    dnums = lax.GatherDimensionNumbers(
        offset_dims=tuple(
            range(
                q_ndim, len(a_shape) + q_ndim if keepdims else len(a_shape) + q_ndim - 1
            )
        ),
        collapsed_slice_dims=() if keepdims else (axis,),
        start_index_map=(axis,),
    )
    low = low[..., None]
    high = high[..., None]
    low_value = lax.gather(a, low, dimension_numbers=dnums, slice_sizes=slice_sizes)
    high_value = lax.gather(a, high, dimension_numbers=dnums, slice_sizes=slice_sizes)
    if q_ndim == 1:
        low_weight = lax.broadcast_in_dim(
            low_weight, low_value.shape, broadcast_dimensions=(0,)
        )
        high_weight = lax.broadcast_in_dim(
            high_weight, high_value.shape, broadcast_dimensions=(0,)
        )

    if interpolation == "linear":
        result = lax.add(
            lax.mul(low_value.astype(q.dtype), low_weight),
            lax.mul(high_value.astype(q.dtype), high_weight),
        )
    elif interpolation == "lower":
        result = low_value
    elif interpolation == "higher":
        result = high_value
    elif interpolation == "nearest":
        pred = lax.le(high_weight, _constant_like(high_weight, 0.5))
        result = lax.select(pred, low_value, high_value)
    elif interpolation == "midpoint":
        result = lax.mul(lax.add(low_value, high_value), _constant_like(low_value, 0.5))
    else:
        raise ValueError(f"interpolation={interpolation!r} not recognized")

    return lax.convert_element_type(result, a.dtype)


# Traceback (TypeError)
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

# Fixed code
def _quantile(a, q, axis, interpolation, keepdims):
    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
    elif isinstance(axis, tuple):
        raise NotImplementedError("Tuple values for axis are not implemented")
    else:
        axis = _canonicalize_axis(axis, ndim(a))

    q_ndim = ndim(q)
    if q_ndim > 1:
        raise ValueError("q must be have rank <= 1, got shape {}".format(shape(q)))

    a_shape = shape(a)
    a = lax.sort(a, dimension=axis)

    n = a_shape[axis]
    q = lax.mul(q, _constant_like(q, n - 1))
    low = lax.floor(q)
    high = lax.ceil(q)
    high_weight = lax.sub(q, low)
    low_weight = lax.sub(_constant_like(high_weight, 1), high_weight)

    low = lax.clamp(_constant_like(low, 0), low, _constant_like(low, n - 1))
    high = lax.clamp(_constant_like(high, 0), high, _constant_like(high, n - 1))
    low = lax.convert_element_type(low, int64)
    high = lax.convert_element_type(high, int64)

    slice_sizes = list(a_shape)
    slice_sizes[axis] = 1

    dnums = lax.GatherDimensionNumbers(
        offset_dims=tuple(
            range(
                q_ndim, len(a_shape) + q_ndim if keepdims else len(a_shape) + q_ndim - 1
            )
        ),
        collapsed_slice_dims=() if keepdims else (axis,),
        start_index_map=(axis,),
    )
    low = low[..., None]
    high = high[..., None]
    low_value = lax.gather(a, low, dimension_numbers=dnums, slice_sizes=slice_sizes)
    high_value = lax.gather(a, high, dimension_numbers=dnums, slice_sizes=slice_sizes)
    if q_ndim == 1:
        low_weight = lax.broadcast_in_dim(
            low_weight, low_value.shape, broadcast_dimensions=(0,)
        )
        high_weight = lax.broadcast_in_dim(
            high_weight, high_value.shape, broadcast_dimensions=(0,)
        )

    if interpolation == "linear":
        result = lax.add(
            lax.mul(low_value.astype(q.dtype), low_weight),
            lax.mul(high_value.astype(q.dtype), high_weight),
        )
    elif interpolation == "lower":
        result = low_value
    elif interpolation == "higher":
        result = high_value
    elif interpolation == "nearest":
        pred = lax.le(high_weight, _constant_like(high_weight, 0.5))
        result = lax.select(pred, low_value, high_value)
    elif interpolation == "midpoint":
        result = lax.mul(lax.add(low_value, high_value), _constant_like(low_value, 0.5))
    else:
        raise ValueError(f"interpolation={interpolation!r} not recognized")

    return lax.convert_element_type(result, a.dtype)