Issue 1:
key_1 = jax.random.PRNGKey(1)
key_2 = jax.random.PRNGKey(2)
print("output for key_1:", jax.random.randint(key_1, (), 0, 2**31))
print("output for key_2:", jax.random.randint(key_2, (), 0, 2**31))

yields

output for key_1: 0
output for key_2: 0

And indeed, the function returns 0 for any key you pass it, so long as maxval >= 2**31.

I guess this is an issue with the dtype, since maxval=2**31-1 works as expected.

Issue 2:
Check for duplicate issues.

Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this:

When run on colab w/ GPU runtime the following code (notebook),

import jax
import numpy as np

print(jax.__version__)

jax.scipy.linalg.sqrtm(np.random.uniform(size=(8,8)))

results in,

Include stack trace

0.3.4
---------------------------------------------------------------------------
JaxStackTraceBeforeTransformation Traceback (most recent call last)
[/usr/lib/python3.7/runpy.py](https://localhost:8080/#) in _run_module_as_main(***failed resolving arguments***)
192 return _run_code(code, main_globals, None,
--> 193 "__main__", mod_spec)
194

53 frames
[/usr/lib/python3.7/runpy.py](https://localhost:8080/#) in _run_code(***failed resolving arguments***)
84 __spec__ = mod_spec)
---> 85 exec(code, run_globals)
86 return run_globals

[/usr/local/lib/python3.7/dist-packages/ipykernel_launcher.py](https://localhost:8080/#) in <module>()
15 from ipykernel import kernelapp as app
---> 16 app.launch_new_instance()

[/usr/local/lib/python3.7/dist-packages/traitlets/config/application.py](https://localhost:8080/#) in launch_instance(***failed resolving arguments***)
845 app.initialize(argv)
--> 846 app.start()
847

[/usr/local/lib/python3.7/dist-packages/ipykernel/kernelapp.py](https://localhost:8080/#) in start(***failed resolving arguments***)
498 try:
--> 499 self.io_loop.start()
500 except KeyboardInterrupt:

[/usr/local/lib/python3.7/dist-packages/tornado/platform/asyncio.py](https://localhost:8080/#) in start(***failed resolving arguments***)
131 asyncio.set_event_loop(self.asyncio_loop)
--> 132 self.asyncio_loop.run_forever()
133 finally:

[/usr/lib/python3.7/asyncio/base_events.py](https://localhost:8080/#) in run_forever(***failed resolving arguments***)
540 while True:
--> 541 self._run_once()
542 if self._stopping:

[/usr/lib/python3.7/asyncio/base_events.py](https://localhost:8080/#) in _run_once(***failed resolving arguments***)
1785 else:
-> 1786 handle._run()
1787 handle = None # Needed to break cycles when an exception occurs.

[/usr/lib/python3.7/asyncio/events.py](https://localhost:8080/#) in _run(***failed resolving arguments***)
87 try:
---> 88 self._context.run(self._callback, *self._args)
89 except Exception as exc:

[/usr/local/lib/python3.7/dist-packages/tornado/platform/asyncio.py](https://localhost:8080/#) in _handle_events(***failed resolving arguments***)
121 fileobj, handler_func = self.handlers[fd]
--> 122 handler_func(fileobj, events)
123

[/usr/local/lib/python3.7/dist-packages/tornado/stack_context.py](https://localhost:8080/#) in null_wrapper(***failed resolving arguments***)
299 _state.contexts = cap_contexts[0]
--> 300 return fn(*args, **kwargs)
301 finally:

[/usr/local/lib/python3.7/dist-packages/zmq/eventloop/zmqstream.py](https://localhost:8080/#) in _handle_events(***failed resolving arguments***)
451 if zmq_events & zmq.POLLIN and self.receiving():
--> 452 self._handle_recv()
453 if not self.socket:

[/usr/local/lib/python3.7/dist-packages/zmq/eventloop/zmqstream.py](https://localhost:8080/#) in _handle_recv(***failed resolving arguments***)
480 callback = self._recv_callback
--> 481 self._run_callback(callback, msg)
482

[/usr/local/lib/python3.7/dist-packages/zmq/eventloop/zmqstream.py](https://localhost:8080/#) in _run_callback(***failed resolving arguments***)
430 # inside our blanket exception handler rather than outside.
--> 431 callback(*args, **kwargs)
432 except Exception:

[/usr/local/lib/python3.7/dist-packages/tornado/stack_context.py](https://localhost:8080/#) in null_wrapper(***failed resolving arguments***)
299 _state.contexts = cap_contexts[0]
--> 300 return fn(*args, **kwargs)
301 finally:

[/usr/local/lib/python3.7/dist-packages/ipykernel/kernelbase.py](https://localhost:8080/#) in dispatcher(***failed resolving arguments***)
282 def dispatcher(msg):
--> 283 return self.dispatch_shell(stream, msg)
284 return dispatcher

[/usr/local/lib/python3.7/dist-packages/ipykernel/kernelbase.py](https://localhost:8080/#) in dispatch_shell(***failed resolving arguments***)
232 try:
--> 233 handler(stream, idents, msg)
234 except Exception:

[/usr/local/lib/python3.7/dist-packages/ipykernel/kernelbase.py](https://localhost:8080/#) in execute_request(***failed resolving arguments***)
398 reply_content = self.do_execute(code, silent, store_history,
--> 399 user_expressions, allow_stdin)
400

[/usr/local/lib/python3.7/dist-packages/ipykernel/ipkernel.py](https://localhost:8080/#) in do_execute(***failed resolving arguments***)
207 try:
--> 208 res = shell.run_cell(code, store_history=store_history, silent=silent)
209 finally:

[/usr/local/lib/python3.7/dist-packages/ipykernel/zmqshell.py](https://localhost:8080/#) in run_cell(***failed resolving arguments***)
536 self._last_traceback = None
--> 537 return super(ZMQInteractiveShell, self).run_cell(*args, **kwargs)
538

[/usr/local/lib/python3.7/dist-packages/IPython/core/interactiveshell.py](https://localhost:8080/#) in run_cell(***failed resolving arguments***)
2717 has_raised = self.run_ast_nodes(code_ast.body, cell_name,
-> 2718 interactivity=interactivity, compiler=compiler, result=result)
2719

[/usr/local/lib/python3.7/dist-packages/IPython/core/interactiveshell.py](https://localhost:8080/#) in run_ast_nodes(***failed resolving arguments***)
2827 code = compiler(mod, cell_name, "single")
-> 2828 if self.run_code(code, result):
2829 return True

[/usr/local/lib/python3.7/dist-packages/IPython/core/interactiveshell.py](https://localhost:8080/#) in run_code(***failed resolving arguments***)
2881 #rprint('Running code', repr(code_obj)) # dbg
-> 2882 exec(code_obj, self.user_global_ns, self.user_ns)
2883 finally:

[<ipython-input-1-f00b9a3dc36a>](https://localhost:8080/#) in <module>()
5
----> 6 jax.scipy.linalg.sqrtm(np.random.uniform(size=(16,16)))

[/usr/local/lib/python3.7/dist-packages/jax/_src/scipy/linalg.py](https://localhost:8080/#) in sqrtm(***failed resolving arguments***)
657 raise NotImplementedError("Blocked version is not implemented yet.")
--> 658 return _sqrtm(A)

[/usr/local/lib/python3.7/dist-packages/jax/_src/scipy/linalg.py](https://localhost:8080/#) in _sqrtm(***failed resolving arguments***)
637 def _sqrtm(A):
--> 638 T, Z = schur(A, output='complex')
639 sqrt_T = _sqrtm_triu(T)

[/usr/local/lib/python3.7/dist-packages/jax/_src/scipy/linalg.py](https://localhost:8080/#) in schur(***failed resolving arguments***)
116 "Expected 'output' to be either 'real' or 'complex', got output={}.".format(output))
--> 117 return _schur(a, output)
118

[/usr/local/lib/python3.7/dist-packages/jax/_src/scipy/linalg.py](https://localhost:8080/#) in _schur(***failed resolving arguments***)
109 a = a.astype(jnp.result_type(a.dtype, 0j))
--> 110 return lax_linalg.schur(a)
111

JaxStackTraceBeforeTransformation: TypeError: _schur_translation_rule() got an unexpected keyword argument 'select_callable'

The preceding stack trace is the source of the JAX operation that, once transformed by JAX, triggered the following exception.

--------------------

The above exception was the direct cause of the following exception:

UnfilteredStackTrace Traceback (most recent call last)
[<ipython-input-1-f00b9a3dc36a>](https://localhost:8080/#) in <module>()
5
----> 6 jax.scipy.linalg.sqrtm(np.random.uniform(size=(16,16)))

[/usr/local/lib/python3.7/dist-packages/jax/_src/scipy/linalg.py](https://localhost:8080/#) in sqrtm(A, blocksize)
657 raise NotImplementedError("Blocked version is not implemented yet.")
--> 658 return _sqrtm(A)

[/usr/local/lib/python3.7/dist-packages/jax/_src/traceback_util.py](https://localhost:8080/#) in reraise_with_filtered_traceback(*args, **kwargs)
161 try:
--> 162 return fun(*args, **kwargs)
163 except Exception as e:

[/usr/local/lib/python3.7/dist-packages/jax/_src/api.py](https://localhost:8080/#) in cache_miss(*args, **kwargs)
434 device=device, backend=backend, name=flat_fun.__name__,
--> 435 donated_invars=donated_invars, inline=inline)
436 out_pytree_def = out_tree()

[/usr/local/lib/python3.7/dist-packages/jax/core.py](https://localhost:8080/#) in bind(self, fun, *args, **params)
1708 def bind(self, fun, *args, **params):
-> 1709 return call_bind(self, fun, *args, **params)
1710

[/usr/local/lib/python3.7/dist-packages/jax/core.py](https://localhost:8080/#) in call_bind(primitive, fun, *args, **params)
1720 tracers = map(top_trace.full_raise, args)
-> 1721 outs = top_trace.process_call(primitive, fun, tracers, params)
1722 return map(full_lower, apply_todos(env_trace_todo(), outs))

[/usr/local/lib/python3.7/dist-packages/jax/core.py](https://localhost:8080/#) in process_call(self, primitive, f, tracers, params)
613 def process_call(self, primitive, f, tracers, params):
--> 614 return primitive.impl(f, *tracers, **params)
615 process_map = process_call

[/usr/local/lib/python3.7/dist-packages/jax/_src/dispatch.py](https://localhost:8080/#) in _xla_call_impl(***failed resolving arguments***)
142 compiled_fun = _xla_callable(fun, device, backend, name, donated_invars,
--> 143 *unsafe_map(arg_spec, args))
144 try:

[/usr/local/lib/python3.7/dist-packages/jax/linear_util.py](https://localhost:8080/#) in memoized_fun(fun, *args)
271 else:
--> 272 ans = call(fun, *args)
273 cache[key] = (ans, fun.stores)

[/usr/local/lib/python3.7/dist-packages/jax/_src/dispatch.py](https://localhost:8080/#) in _xla_callable_uncached(fun, device, backend, name, donated_invars, *arg_specs)
169 return lower_xla_callable(fun, device, backend, name, donated_invars,
--> 170 *arg_specs).compile().unsafe_call
171

[/usr/local/lib/python3.7/dist-packages/jax/_src/profiler.py](https://localhost:8080/#) in wrapper(*args, **kwargs)
205 with TraceAnnotation(name, **decorator_kwargs):
--> 206 return func(*args, **kwargs)
207 return wrapper

[/usr/local/lib/python3.7/dist-packages/jax/_src/dispatch.py](https://localhost:8080/#) in lower_xla_callable(fun, device, backend, name, donated_invars, *arg_specs)
259 module_name, closed_jaxpr, backend.platform,
--> 260 mlir.ReplicaAxisContext(axis_env), name_stack, donated_invars)
261 else:

[/usr/local/lib/python3.7/dist-packages/jax/interpreters/mlir.py](https://localhost:8080/#) in lower_jaxpr_to_module(module_name, jaxpr, platform, axis_context, name_stack, donated_args, replicated_args, arg_shardings, result_shardings)
493 arg_shardings=arg_shardings, result_shardings=result_shardings,
--> 494 input_output_aliases=input_output_aliases)
495

[/usr/local/lib/python3.7/dist-packages/jax/interpreters/mlir.py](https://localhost:8080/#) in lower_jaxpr_to_fun(ctx, name, jaxpr, public, replace_units_with_dummy, replace_tokens_with_dummy, replicated_args, arg_shardings, result_shardings, input_output_aliases)
636 jaxpr.jaxpr, map(ir_constants, jaxpr.consts),
--> 637 *args)
638 outs = []

[/usr/local/lib/python3.7/dist-packages/jax/interpreters/mlir.py](https://localhost:8080/#) in jaxpr_subcomp(ctx, jaxpr, consts, *args)
722 ans = rule(rule_ctx, *map(_unwrap_singleton_ir_values, in_nodes),
--> 723 **eqn.params)
724

[/usr/local/lib/python3.7/dist-packages/jax/interpreters/mlir.py](https://localhost:8080/#) in _xla_call_lower(***failed resolving arguments***)
783 backend, ctx.module_context, ctx.avals_in, ctx.avals_out,
--> 784 *args)
785

[/usr/local/lib/python3.7/dist-packages/jax/interpreters/mlir.py](https://localhost:8080/#) in _call_lowering(fn_name, stack_name, call_jaxpr, backend, ctx, avals_in, avals_out, *args)
771 symbol_name = lower_jaxpr_to_fun(sub_ctx, fn_name,
--> 772 core.ClosedJaxpr(call_jaxpr, ())).name.value
773 call = func_dialect.CallOp(flat_output_types,

[/usr/local/lib/python3.7/dist-packages/jax/interpreters/mlir.py](https://localhost:8080/#) in lower_jaxpr_to_fun(ctx, name, jaxpr, public, replace_units_with_dummy, replace_tokens_with_dummy, replicated_args, arg_shardings, result_shardings, input_output_aliases)
636 jaxpr.jaxpr, map(ir_constants, jaxpr.consts),
--> 637 *args)
638 outs = []

[/usr/local/lib/python3.7/dist-packages/jax/interpreters/mlir.py](https://localhost:8080/#) in jaxpr_subcomp(ctx, jaxpr, consts, *args)
722 ans = rule(rule_ctx, *map(_unwrap_singleton_ir_values, in_nodes),
--> 723 **eqn.params)
724

[/usr/local/lib/python3.7/dist-packages/jax/interpreters/mlir.py](https://localhost:8080/#) in cached_lowering(ctx, *args, **params)
935 if func is None:
--> 936 func = _emit_lowering_rule_as_fun(partial(f, **params), ctx)
937 ctx.module_context.cached_primitive_lowerings[key] = func

[/usr/local/lib/python3.7/dist-packages/jax/interpreters/mlir.py](https://localhost:8080/#) in _emit_lowering_rule_as_fun(lowering_rule, ctx)
665 map(len, input_types))
--> 666 outs = lowering_rule(ctx, *_unwrap_singleton_ir_values(unflattened_args))
667 func_dialect.ReturnOp(util.flatten(map(wrap_singleton_ir_values, outs)))

[/usr/local/lib/python3.7/dist-packages/jax/interpreters/mlir.py](https://localhost:8080/#) in fallback(ctx, *args, **params)
952 xla_computation = xla.primitive_subcomputation(
--> 953 module_ctx.platform, module_ctx.axis_env, prim, *ctx.avals_in, **params)
954 submodule_str = xc._xla.mlir.xla_computation_to_mlir_module(xla_computation)

[/usr/local/lib/python3.7/dist-packages/jax/interpreters/xla.py](https://localhost:8080/#) in primitive_subcomputation(platform, axis_env, prim, *avals, **params)
445 name_stack=new_name_stack())
--> 446 ans = f(ctx.replace(builder=c), avals, None, *xla_args, **params)
447 if prim.multiple_results:

[/usr/local/lib/python3.7/dist-packages/jax/interpreters/xla.py](https://localhost:8080/#) in f_new(ctx, avals_in, avals_out, *xla_args, **params)
1035 return jaxpr_subcomp(ctx, jaxpr, _xla_consts(ctx.builder, consts),
-> 1036 *xla_args)
1037 return f_new

[/usr/local/lib/python3.7/dist-packages/jax/interpreters/xla.py](https://localhost:8080/#) in jaxpr_subcomp(ctx, jaxpr, consts, *args)
611 ans = rule(eqn_ctx, map(aval, eqn.invars), map(aval, eqn.outvars),
--> 612 *in_nodes, **eqn.params)
613

UnfilteredStackTrace: TypeError: _schur_translation_rule() got an unexpected keyword argument 'select_callable'

The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.

--------------------

The above exception was the direct cause of the following exception:

TypeError Traceback (most recent call last)
[<ipython-input-1-f00b9a3dc36a>](https://localhost:8080/#) in <module>()
4 print(jax.__version__)
5
----> 6 jax.scipy.linalg.sqrtm(np.random.uniform(size=(16,16)))

[/usr/local/lib/python3.7/dist-packages/jax/_src/scipy/linalg.py](https://localhost:8080/#) in sqrtm(A, blocksize)
656 if blocksize > 1:
657 raise NotImplementedError("Blocked version is not implemented yet.")
--> 658 return _sqrtm(A)

TypeError: _schur_translation_rule() got an unexpected keyword argument 'select_callable'