Ipython and jupyter notebook keep failing while testing jax

[Sam131112]

python version 3.6.10
cuda version 11
jax version '0.2.17'
gcc 7.3.0

import jax.numpy as jnp
from jax import grad, jit, vmap
from jax import random
import jax
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
key = random.PRNGKey(0)
size = 100
x = random.normal(key, (size, size), dtype=jnp.float32)
%timeit jnp.dot(x, x.T).block_until_ready()  # runs on the GPU

Error obtained

2021-07-10 19:06:52.914859: W external/org_tensorflow/tensorflow/stream_executor/gpu/asm_compiler.cc:77] Couldn't get ptxas version string: Internal: Couldn't invoke ptxas --version
2021-07-10 19:06:52.915444: F external/org_tensorflow/tensorflow/compiler/xla/service/gpu/nvptx_compiler.cc:480] ptxas returned an error during compilation of ptx to sass: 'Internal: Failed to launch ptxas' If the error message indicates that a file could not be written, please verify that sufficient filesystem space is provided.
Aborted (core dumped)

[hawkinsp]

This means that JAX could not find ptxas, which is part of the CUDA toolkit. Is it in your path? If it's in your PATH jax should find it.

[oceanumeric]

The same thing happened to me when I was running the following code in Jupyter (with VScode), whereas the code works well in the terminal:

# test random from jax 
key = jrd.PRNGKey(0)
x = jrd.normal(key, (10,))
print(x)

Here is the error:

2022-12-03 10:15:18.889345: E external/org_tensorflow/tensorflow/compiler/xla/stream_executor/gpu/asm_compiler.cc:114] *** WARNING *** You are using ptxas 10.1.243, which is older than 11.1. ptxas before 11.1 is known to miscompile XLA code, leading to incorrect results or invalid-address errors.

2022-12-03 10:15:18.890789: W external/org_tensorflow/tensorflow/compiler/xla/stream_executor/gpu/asm_compiler.cc:231] Falling back to the CUDA driver for PTX compilation; ptxas does not support CC 8.6
2022-12-03 10:15:18.890811: W external/org_tensorflow/tensorflow/compiler/xla/stream_executor/gpu/asm_compiler.cc:234] Used ptxas at ptxas
2022-12-03 10:15:18.892319: E external/org_tensorflow/tensorflow/compiler/xla/stream_executor/cuda/cuda_driver.cc:628] failed to get PTX kernel "shift_right_logical" from module: CUDA_ERROR_NOT_FOUND: named symbol not found
2022-12-03 10:15:18.892343: E external/org_tensorflow/tensorflow/compiler/xla/pjrt/pjrt_stream_executor_client.cc:2153] Execution of replica 0 failed: INTERNAL: Could not find the corresponding function

No idea why.

The weird thing is that it runs in the jupyter when you open the file via the command line jupyter notebook.

[jakevdp]

The python environment used by Jupyter, by VSCode, and by your command-line are not necessarily the same. I'd run this in all cases:

import sys
print(sys.executable)

I suspect the executables are different, and that JAX is incorrectly installed for the executable used by VSCode.

[oceanumeric]

Thanks for your reply @jakevdp

I just ran print(sys.executable), and the returned paths are the same.

I think it is because of VScode using different environment path and JAX cannot search for the right cuda.

[oceanumeric]

Solve it with Docker Container, here is how you can set it up: https://github.com/oceanumeric/cuda-container-vscode

[oceanumeric]

There is another way to make it work: within virtual environment,

• launch jupyter notebook
• copy jupyter server link
• On the right corner of VS Code click Jupyter Server
• paste the server link and switch the kernel to the one you created before (for instance, python3 -m ipykernel install --user --name=myvenv

This will make sure everything running will be within the virtual environment (this avoids the case that VS Code Jupyter extension is installed globally when you launch the project).