Fall back to CPU for jax.scipy.linalg.schur?

[nalzok]

It is known that when using a GPU or TPU backend, calling jax.scipy.linalg.schur raises NotImplementedError: Schur decomposition is only implemented on the CPU backend. I wonder if it is possible to bypass the issue by falling back to the CPU implementation in such cases.

I tried to transfer the array back to the CPU with jax.device_get(X), hoping that would trigger the LAPACK backend, but it still gives the same NotImplementedError.

import jax
X = jax.numpy.array([[1, 2], [3, 4]])
jax.scipy.linalg.schur(jax.device_get(X))

To clarify, I only want this particular function to be executed on CPU, since there isn't a GPU/TPU implementation. The idea is that I will do a lot of computation on a GPU/TPU, transfer the array to the CPU, perform Schur decomposition, transfer the result back to GPT/TPU, and continue execution there.

[sharadmv]

You could use jax.pure_callback to call out to scipy.linalg.schur.

import jax
import jax.numpy as jnp
import scipy.linalg

def schur(x):
  return jax.pure_callback(scipy.linalg.schur, (x, x), x)

@jax.jit
def f(x):
  return schur(x)

print(f(jnp.array([[0, 2, 2], [0, 1, 2], [1, 0, 1]], jnp.float32)))

[nalzok]

This doesn't really work; see the stack trace below. Does it work on your JAX installation?

Python 3.8.10 (default, Jun 22 2022, 20:18:18)
[GCC 9.4.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import jax
>>> import jax.numpy as jnp
>>> import scipy.linalg
>>>
>>> def schur(x):
...   return jax.pure_callback(scipy.linalg.schur, (x, x), x)
...
>>> @jax.jit
... def f(x):
...   return schur(x)
...
>>> print(f(jnp.array([[0, 2, 2], [0, 1, 2], [1, 0, 1]], jnp.float32)))
F0907 14:29:33.631891  939012 host_command_dispatcher.cc:83] Check failed: !handlers_by_run_ids_[queue_id].empty() Host command 50331929 triggered but no handler was registered, run id: 1
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/qys/Research/label-shift/.venv/lib/python3.8/site-packages/jax/_src/traceback_util.py", line 162, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/home/qys/Research/label-shift/.venv/lib/python3.8/site-packages/jax/_src/api.py", line 528, in cache_miss
    out_flat = xla.xla_call(
  File "/home/qys/Research/label-shift/.venv/lib/python3.8/site-packages/jax/core.py", line 1963, in bind
    return call_bind(self, fun, *args, **params)
  File "/home/qys/Research/label-shift/.venv/lib/python3.8/site-packages/jax/core.py", line 1979, in call_bind
    outs = top_trace.process_call(primitive, fun_, tracers, params)
  File "/home/qys/Research/label-shift/.venv/lib/python3.8/site-packages/jax/core.py", line 689, in process_call
    return primitive.impl(f, *tracers, **params)
  File "/home/qys/Research/label-shift/.venv/lib/python3.8/site-packages/jax/_src/dispatch.py", line 236, in _xla_call_impl
    return compiled_fun(*args)
  File "/home/qys/Research/label-shift/.venv/lib/python3.8/site-packages/jax/_src/dispatch.py", line 841, in _execute_compiled
    out_bufs = token_handler(out_bufs, runtime_token)
  File "/home/qys/Research/label-shift/.venv/lib/python3.8/site-packages/jax/_src/dispatch.py", line 808, in _remove_tokens
    output_token_buf, *token_bufs = token_bufs
jax._src.traceback_util.UnfilteredStackTrace: ValueError: not enough values to unpack (expected at least 1, got 0)

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:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
ValueError: not enough values to unpack (expected at least 1, got 0)
>>> *** Check failure stack trace: ***
    @     0x7f07a33b98e4  (unknown)
    @     0x7f07a33b93ca  (unknown)
    @     0x7f07a33b9c49  (unknown)
    @     0x7f079ff89e79  (unknown)
    @     0x7f079ff70f66  (unknown)
    @     0x7f079dfa17c3  (unknown)
    @     0x7f07a3297669  (unknown)
    @     0x7f07a32930fe  (unknown)
    @     0x7f0e8753b609  start_thread
https://symbolize.stripped_domain/r/?trace=7f07a33b98e4,7f07a33b93c9,7f07a33b9c48,7f079ff89e78,7f079ff70f65,7f079dfa17c2,7f07a3297668,7f07a32930fd,7f0e8753b608&map=068bf80b76f830987166dd8847d0248f:7f078dddc000-7f07a370ede0
https://symbolize.stripped_domain/r/?trace=7f0e8759900b,7f0e8759908f,7f07a33b9943,7f07a33b93c9,7f07a33b9c48,7f079ff89e78,7f079ff70f65,7f079dfa17c2,7f07a3297668,7f07a32930fd,7f0e8753b608&map=068bf80b76f830987166dd8847d0248f:7f078dddc000-7f07a370ede0
*** SIGABRT received by PID 937881 (TID 939012) on cpu 81 from PID 937881; ***
E0907 14:29:33.636138  939012 coredump_hook.cc:370] RAW: Remote crash data gathering hook invoked.
E0907 14:29:33.636153  939012 coredump_hook.cc:416] RAW: Skipping coredump since rlimit was 0 at process start.
E0907 14:29:33.636161  939012 client.cc:242] RAW: Coroner client retries enabled (b/136286901), will retry for up to 30 sec.
E0907 14:29:33.636166  939012 coredump_hook.cc:477] RAW: Sending fingerprint to remote end.
E0907 14:29:33.636174  939012 coredump_socket.cc:118] RAW: Stat failed errno=2 on socket /var/google/services/logmanagerd/remote_coredump.socket
E0907 14:29:33.636184  939012 coredump_hook.cc:481] RAW: Cannot send fingerprint to Coroner: [NOT_FOUND] Missing crash reporting socket. Is the listener running?
E0907 14:29:33.636190  939012 coredump_hook.cc:555] RAW: Discarding core.
F0907 14:29:33.631891  939012 host_command_dispatcher.cc:83] Check failed: !handlers_by_run_ids_[queue_id].empty() Host command 50331929 triggered but no handler was registered, run id: 1
E0907 14:29:33.898010  939012 process_state.cc:774] RAW: Raising signal 6 with default behavior
fish: “pipenv run python3” terminated by signal SIGABRT (Abort)

[sharadmv]

Could you file an issue w/ JAX+jaxlib versions and platform info and we can look into it. This should work cross-platform.

[nalzok]

Sure. It's here: #12260

[sharadmv]

Thank you, will investigate