WARNING: jax 0.4.19 does not provide the extra 'cuda12-pip'

[kwang0913]

Hi guys,

I received this warring

WARNING: jax 0.4.19 does not provide the extra 'cuda12-pip'

during installation of the current jax library. It is still installed successfully. I'm just wondering what would be the correct command to install jax without getting this warning. Currently I install the library by

pip install -U "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

Besides, I'm getting the following error and warning after installation

2023-10-22 09:33:14.822791: E tensorflow/compiler/xla/stream_executor/cuda/cuda_dnn.cc:9342] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered 2023-10-22 09:33:14.822814: E tensorflow/compiler/xla/stream_executor/cuda/cuda_fft.cc:609] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered 2023-10-22 09:33:14.822832: E tensorflow/compiler/xla/stream_executor/cuda/cuda_blas.cc:1518] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered 2023-10-22 09:33:15.459916: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT 2023-10-22 09:33:16.073033: W external/xla/xla/service/gpu/nvptx_compiler.cc:673] The NVIDIA driver's CUDA version is 12.2 which is older than the ptxas CUDA version (12.3.52). Because the driver is older than the ptxas version, XLA is disabling parallel compilation, which may slow down compilation. You should update your NVIDIA driver or use the NVIDIA-provided CUDA forward compatibility packages.

Thank you for any help.

Best

[jakevdp]

I think you're looking for cuda12_pip, not cuda12-pip