Repository: google/jax

Issue:
[GPU] jax.scipy.linalg.sqrtm --> TypeError: _schur_translation_rule() got an unexpected keyword argument 'select_callable'
- [x] Check for duplicate issues.
- [x] 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](https://colab.research.google.com/drive/1zFK5h_6Jnzw5kPRKpR9Dko4kVpQa2qxL)),
```
import jax
import numpy as np

print(jax.__version__)

jax.scipy.linalg.sqrtm(np.random.uniform(size=(8,8)))
```
results in,

- [x] 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'
```



Discussion:

Resolving patch:
diff --git a/jax/_src/lax/linalg.py b/jax/_src/lax/linalg.py
--- a/jax/_src/lax/linalg.py
+++ b/jax/_src/lax/linalg.py
@@ -1816,8 +1816,7 @@ def _schur_impl(operand, *, compute_schur_vectors, sort_eig_vals,
       select_callable=select_callable)
 
 
-def _schur_translation_rule(ctx, avals_in, avals_out, operand, *,
-                            compute_schur_vectors, sort_eig_vals):
+def _schur_translation_rule(ctx, *args, **kwargs):
   raise NotImplementedError(
       "Schur decomposition is only implemented on the CPU backend.")
 


Test patch: