Unable to use jax

[Haislich]

I installed Jax as per instructions

However this simple code

from jax import numpy as jnp

a = jnp.array([0] * (3 * 210 * 160))
a = a.reshape((3, 210, 160))
a = jnp.resize(a, (1, 110, 84))

Causes this issue

Traceback (most recent call last):
  File "/home/haislich/Documents/dqn/src/dqn/dqn.py", line 5, in <module>
    a = jnp.resize(a, (1, 110, 84))
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^
jaxlib.xla_extension.XlaRuntimeError: INTERNAL: ptxas exited with non-zero error code 139, output: ptxas /tmp/tempfile-Workstation-c7212f129f34c799-9923-629b5d77da91f, line 5; fatal   : Unsupported .version 8.3; current version is '8.0'

--------------------
For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

What can I do ?

[jakevdp]

Hi - sorry you're having this issue!

Can you give us some more information about what system you're on, what Python version you have, how you installed JAX, and what JAX version you're using?

[Haislich]

What system you're on

lsb_release -a

Distributor ID:	Ubuntu
Description:	Ubuntu 24.04.1 LTS
Release:	24.04
Codename:	noble

uname -a

Linux Workstation 6.8.0-51-generic #52-Ubuntu SMP PREEMPT_DYNAMIC Thu Dec  5 13:09:44 UTC 2024 x86_64 x86_64 x86_64 GNU/Linux

nvidia-smi --version

NVIDIA-SMI version  : 560.35.05
NVML version        : 560.35
DRIVER version      : 560.35.05
CUDA Version        : 12.6

nvcc --version

nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2023 NVIDIA Corporation
Built on Fri_Jan__6_16:45:21_PST_2023
Cuda compilation tools, release 12.0, V12.0.140
Build cuda_12.0.r12.0/compiler.32267302_0

what Python version you have

python --version

Python 3.12.7

how you installed JAX

pip install --upgrade "jax[cuda12]"

what JAX version you're using

pip freeze | grep jax

jax==0.4.38
jax-cuda12-pjrt==0.4.38
jax-cuda12-plugin==0.4.38
jaxlib==0.4.38

[jakevdp]

Thanks – what's your CUDNN version?

[Haislich]

Isn't the CUDNN version just the result of nvcc --version ?

[pearu]

This kind of an error may appear when cuda toolkit and nvidia driver are incompatible: nvcc reports 12.0 but the driver is 12.6. Check also the output of

apt list --installed | grep cuda

and consider updating cuda toolkit.

[Haislich]

I tried updating it, however nothing happened as resulted that the latest version was installed.
However you diagnosed the problem correctly there was a mismatch in drivers, I uninstalled all non 12.6 version that I found and rebooted and now it works.
Thank you very much for the help!