Issue 1:
Description

jnp.arange causes XlaRuntimeError.

!pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

!nvidia-smi

import jax.numpy as jnp

x = jnp.arange(33)

XlaRuntimeError Traceback (most recent call last)

in

----> 1 x = jnp.arange(33)

16 frames

/usr/local/lib/python3.8/dist-packages/jax/_src/dispatch.py in backend_compile(backend, built_c, options, host_callbacks)

1024 # TODO(sharadmv): remove this fallback when all backends allow compile

1025 # to take in host_callbacks

-> 1026 return backend.compile(built_c, compile_options=options)

1027

1028 _ir_dump_counter = itertools.count()

XlaRuntimeError: INTERNAL: RET_CHECK failure (external/org_tensorflow/tensorflow/compiler/xla/service/gpu/gpu_compiler.cc:626) dnn != nullptr

What jax/jaxlib version are you using?

v0.4.3

Which accelerator(s) are you using?

GPU

Additional system info

Google Colab

NVIDIA GPU info

Tue Feb 14 16:45:09 2023

+-----------------------------------------------------------------------------+

| NVIDIA-SMI 510.47.03 Driver Version: 510.47.03 CUDA Version: 11.6 |

|-------------------------------+----------------------+----------------------+

| GPU Name Persistence-M| Bus-Id Disp.A | Volatile Uncorr. ECC |

| Fan Temp Perf Pwr:Usage/Cap| Memory-Usage | GPU-Util Compute M. |

| | | MIG M. |

|===============================+======================+======================|

| 0 Tesla T4 Off | 00000000:00:04.0 Off | 0 |

| N/A 59C P0 28W / 70W | 0MiB / 15360MiB | 0% Default |

| | | N/A |

+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+

| Processes: |

| GPU GI CI PID Type Process name GPU Memory |

| ID ID Usage |

|=============================================================================|

| No running processes found |

+-----------------------------------------------------------------------------+

Issue 2:
Description

I am getting an XLA weird error when running

rng = jax.random.PRNGKey(0)

in Google colab. You can check the error here.

The full code is

!pip install --upgrade pip
!pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
!pip install flax
!pip install optax

import jax

rng = jax.random.PRNGKey(0)

and returns

---------------------------------------------------------------------------
XlaRuntimeError Traceback (most recent call last)
[<ipython-input-2-71f91d108500>](https://localhost:8080/#) in <module>
----> 1 rng = jax.random.PRNGKey(0)
2 rng, init_rng = jax.random.split(rng)

23 frames
[/usr/local/lib/python3.8/dist-packages/jax/_src/dispatch.py](https://localhost:8080/#) in backend_compile(backend, built_c, options, host_callbacks)
1024 # TODO(sharadmv): remove this fallback when all backends allow `compile`
1025 # to take in `host_callbacks`
-> 1026 return backend.compile(built_c, compile_options=options)
1027
1028 _ir_dump_counter = itertools.count()

XlaRuntimeError: INTERNAL: RET_CHECK failure (external/org_tensorflow/tensorflow/compiler/xla/service/gpu/gpu_compiler.cc:626) dnn != nullptr

Also reported here #14463 (comment)

What jax/jaxlib version are you using?

jax == 0.4.3; jaxlib == 0.4.3+cuda11.cudnn86

Which accelerator(s) are you using?

GPU

Additional system info

Google colab

NVIDIA GPU info

Tue Feb 14 11:19:59 2023
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 510.47.03 Driver Version: 510.47.03 CUDA Version: 11.6 |
|-------------------------------+----------------------+----------------------+
| GPU Name Persistence-M| Bus-Id Disp.A | Volatile Uncorr. ECC |
| Fan Temp Perf Pwr:Usage/Cap| Memory-Usage | GPU-Util Compute M. |
| | | MIG M. |
|===============================+======================+======================|
| 0 Tesla T4 Off | 00000000:00:04.0 Off | 0 |
| N/A 46C P0 26W / 70W | 13962MiB / 15360MiB | 0% Default |
| | | N/A |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes: |
| GPU GI CI PID Type Process name GPU Memory |
| ID ID Usage |
|=============================================================================|
| 0 N/A N/A 1906 C 13959MiB |
+-----------------------------------------------------------------------------+