Check failed: stream->parent()->GetBlasGemmAlgorithms(&algorithms)

[MaxH1996]

I am currently playing around with neural ODEs and I wanted to run the following sample code for given here. It seems that I get an error when running:

@jit
def odenet_update(params, inputs, targets):
  grads = grad(odenet_loss)(params, inputs, targets)
  return [(w - step_size * dw, b - step_size * db)
          for (w, b), (dw, db) in zip(params, grads)]

# Initialize and train ODE-Net.
odenet_params = init_random_params(param_scale, odenet_layer_sizes)

for i in range(train_iters):
  odenet_params = odenet_update(odenet_params, inputs, targets)

and the error seems to arise in the for loop. The message that I am getting is:

2021-09-07 11:48:18.112359: W external/org_tensorflow/tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory
2021-09-07 11:48:18.202979: W external/org_tensorflow/tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory
2021-09-07 11:48:18.204374: W external/org_tensorflow/tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory
2021-09-07 11:48:20.020339: W external/org_tensorflow/tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcublas.so.11'; dlerror: libcublas.so.11: cannot open shared object file: No such file or directory
2021-09-07 11:48:20.020361: E external/org_tensorflow/tensorflow/stream_executor/cuda/cuda_blas.cc:226] failed to create cublas handle: CUBLAS_STATUS_INTERNAL_ERROR
2021-09-07 11:48:20.020365: F external/org_tensorflow/tensorflow/compiler/xla/service/gpu/gemm_algorithm_picker.cc:118] Check failed: stream->parent()->GetBlasGemmAlgorithms(&algorithms) 
Aborted (core dumped)

I have looked at some posts similar to this, but a lot of them give different solutions, or are not entirely relevant to my problem. Any help on this would be appreciated.

[hawkinsp]

That's a particularly ungraceful failure, but the underlying problem seems to be that JAX can't find cudart or cublas, which are parts of the CUDA toolkit. Do you have CUDA installed? If so, which version and where? What version of the jaxlib wheel do you have installed? If you have CUDA installed and a matching jaxlib, you might try adding the lib directory in your CUDA installation to LD_LIBRARY_PATH.

[MaxH1996]

Ah, I think I made a mistake because I believe it doesn't point anywhere. However, I am not quite sure where my cuda path actually is, I assumed it would be /usr/local/cuda but that doesn't exist. When looking for the path using:

locate cuda | grep /cuda$

I get these:

/usr/include/thrust/system/cuda
/usr/lib/cuda
/usr/share/doc/libthrust-dev/examples/cuda

[hawkinsp]

How did you install CUDA?

[MaxH1996]

Using sudo apt and I haven't had problems with it yet tbh. I usually use PyTorch and never ran into any cuda problems before.

[hawkinsp]

I believe PyTorch bundles CUDA with its wheels, which makes them far larger. JAX does not: you need a working CUDA installation.

If you installed CUDA from Nvidia's packages correctly, you should end up with CUDA linked from /usr/local/cuda. You need to fix your CUDA installation. I believe the correct package is called something like cuda-toolkit-11-2, once you have added NVidia's APT sources.

[MaxH1996]

Great! Thanks for you help :)