Can't find ptxas binary in ${CUDA_DIR}/bin.

[ZisongXu]

2024-02-23 16:57:35.344498: W external/xla/xla/service/gpu/nvptx_compiler.cc:499] Can't find ptxas binary in ${CUDA_DIR}/bin.  Custom ptxas location can be specified using $PATH.
Searched for CUDA in the following directories:
  /home/zisongxu/.local/lib/python3.8/site-packages/nvidia/cuda_nvcc
  /usr/local/cuda-12.0
  /usr/local/cuda
  .
You can choose the search directory by setting xla_gpu_cuda_data_dir in HloModule's DebugOptions.  For most apps, setting the environment variable XLA_FLAGS=--xla_gpu_cuda_data_dir=/path/to/cuda will work.
2024-02-23 16:57:35.344529: F external/xla/xla/service/gpu/nvptx_compiler.cc:503] Can't find ptxas binary.  You can pass the flag --xla_gpu_unsafe_fallback_to_driver_on_ptxas_not_found to use the GPU driver for compiling ptx instead. However this option is discouraged and can lead to increased memory consumptions and other subtle runtime issues.
Aborted (core dumped)

Sorry, i am the new player of cuda/Jax, when I run my code I met this problem. Can anyone tell me how to solve this problem? Thanks a lot!

[hawkinsp]

How did you install JAX? What versions of CUDA do you have installed locally?

I'll also note that we no longer ship release for Python 3.8: you should upgrade.

[ZisongXu]

Hi thanks for the replying, I used

pip install --upgrade pip

# CUDA 12 installation
# Note: wheels only available on linux.
pip install --upgrade "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

to install JAX. The Nvidia-driver cuda version is 12.2