Missing lapack_sgesdd when running jax function with jnp.linalg.svd in C++.

[louis-shao]

When I dump a computation with jnp.linalg.svd to hlo module and execute it in C++ according to the example in the repo, I got an error saying

2023-10-11 18:36:52.070030: E xla/service/cpu/simple_orc_jit.cc:421] Unable to resolve runtime symbol: 'lapack_sgesdd'. Hint: if the symbol a custom call target, make sure you've registered it with the JIT using XLA_CPU_REGISTER_CUSTOM_CALL_TARGET. JIT session error: Symbols not found: [ lapack_sgesdd ].

How do I register that?

[hawkinsp]

That's correct. You'd need to link in and register the various LAPACK custom calls JAX uses. We don't have a well-trodden path for doing this at the moment. Fundamentally you need to build these three modules:

• https://github.com/google/jax/blob/main/jaxlib/cpu/cpu_kernels.cc
• https://github.com/google/jax/blob/main/jaxlib/cpu/lapack_kernels_using_lapack.cc
• https://github.com/google/jax/blob/main/jaxlib/cpu/lapack_kernels.cc
  together with a copy of LAPACK and link them into the C++ code.

Can you say a bit more about the use case? I'm wondering how we could do better here. Which example code are you following?

[louis-shao]

The BUILD file under jax/examples/jax_cpp also depends on openxla, so if I build from the jax codebase and add the following rules to my BUILD file, would it just work?

• //jaxlib/cpu:cpu_kernels
• //jaxlib/cpu:lapack_kernels
• //jaxlib/cpu:lapack_kernels_using_lapack

[louis-shao]

nvm, I extracted the subset of code for lapack_sgesdd from the 3 files you mentioned in the jax codebase and put it inside my codebase. Then I can continue using openxla/xla library to build and run just fine.

[joeycarter]

Hi @hawkinsp,

I just came across this discussion, I'm interested if there have been any developments since. I'd like to do something similar, but in a more generic way to be able to support all of the functions under jax.scipy.linalg. We could copy over all of the custom LAPACK wrappers listed in jaxlib/cpu/cpu_kernels.cc, but this would be tedious and difficult to maintain. Is there a way to link directly to these wrappers included in JAX?

For a bit more context, we have our own JIT compiler, called qjit, and we would like to be able to do the following, for example:

import jax
import jax.numpy as jnp
import qjit

@qjit
def matrix_sqrt(A):
    return jax.scipy.linalg.sqrtm(A)

X = jnp.array([[1., 2., 3.],
               [2., 4., 2.],
               [3., 2., 1.]])
matrix_sqrt(X)

and similarly for the other functions in jax.scipy.linalg. However, when we do so, we get a similar error to the one discussed above, in this case:

OSError: /tmp/matrix_sqrtonhryf12/matrix_sqrt.so: undefined symbol: lapack_zgees

since we aren't able to link to the LAPACK library used by JAX (or at least, we aren't aware of a way to do this).

[dfm]

We don't currently have a good way for you to link to the compiled jaxlib directly. Instead, you'll need to compile and link to the sources that @hawkinsp linked to above, along with your own LAPACK library. I'm hoping that we can have a better workflow in the future, but there hasn't been any real progress on that yet!

[joeycarter]

Okay thanks @dfm! I'll give that a try.