<h3 dir="auto">Description</h3>
<p dir="auto">jnp.arange causes XlaRuntimeError.</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="!pip install --upgrade &quot;jax[cuda]&quot; -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html"><pre class="notranslate"><code class="notranslate">!pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
</code></pre></div>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="!nvidia-smi"><pre class="notranslate"><code class="notranslate">!nvidia-smi
</code></pre></div>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="import jax.numpy as jnp"><pre class="notranslate"><code class="notranslate">import jax.numpy as jnp
</code></pre></div>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="x = jnp.arange(33)"><pre class="notranslate"><code class="notranslate">x = jnp.arange(33)
</code></pre></div>
<hr>
<p dir="auto">XlaRuntimeError                           Traceback (most recent call last)</p>
<p dir="auto"><a href="https://localhost:8080/#" rel="nofollow"></a> in <br>
----&gt; 1 x = jnp.arange(33)</p>
<p dir="auto">16 frames</p>
<p dir="auto"><a href="https://localhost:8080/#" rel="nofollow">/usr/local/lib/python3.8/dist-packages/jax/_src/dispatch.py</a> in backend_compile(backend, built_c, options, host_callbacks)<br>
1024   # TODO(sharadmv): remove this fallback when all backends allow <code class="notranslate">compile</code><br>
1025   # to take in <code class="notranslate">host_callbacks</code><br>
-&gt; 1026   return backend.compile(built_c, compile_options=options)<br>
1027<br>
1028 _ir_dump_counter = itertools.count()</p>
<p dir="auto">XlaRuntimeError: INTERNAL: RET_CHECK failure (external/org_tensorflow/tensorflow/compiler/xla/service/gpu/gpu_compiler.cc:626) dnn != nullptr</p>
<h3 dir="auto">What jax/jaxlib version are you using?</h3>
<p dir="auto">v0.4.3</p>
<h3 dir="auto">Which accelerator(s) are you using?</h3>
<p dir="auto">GPU</p>
<h3 dir="auto">Additional system info</h3>
<p dir="auto">Google Colab</p>
<h3 dir="auto">NVIDIA GPU info</h3>
<p dir="auto">Tue Feb 14 16:45:09 2023<br>
+-----------------------------------------------------------------------------+<br>
| NVIDIA-SMI 510.47.03    Driver Version: 510.47.03    CUDA Version: 11.6     |<br>
|-------------------------------+----------------------+----------------------+<br>
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |<br>
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |<br>
|                               |                      |               MIG M. |<br>
|===============================+======================+======================|<br>
|   0  Tesla T4            Off  | 00000000:00:04.0 Off |                    0 |<br>
| N/A   59C    P0    28W /  70W |      0MiB / 15360MiB |      0%      Default |<br>
|                               |                      |                  N/A |<br>
+-------------------------------+----------------------+----------------------+</p>
<p dir="auto">+-----------------------------------------------------------------------------+<br>
| Processes:                                                                  |<br>
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |<br>
|        ID   ID                                                   Usage      |<br>
|=============================================================================|<br>
|  No running processes found                                                 |<br>
+-----------------------------------------------------------------------------+</p>
<h3 dir="auto">Description</h3>
<p dir="auto">I am getting an XLA weird error when running</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="rng = jax.random.PRNGKey(0)"><pre class="notranslate"><span class="pl-s1">rng</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-s1">random</span>.<span class="pl-v">PRNGKey</span>(<span class="pl-c1">0</span>)</pre></div>
<p dir="auto">in Google colab. You can check the error <a href="https://colab.research.google.com/drive/1nAqTAf-cxC1bHeTC8Pfk8v82REf7t38V?usp=sharing" rel="nofollow">here</a>.<br>
The full code is</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="!pip install --upgrade pip
!pip install --upgrade &quot;jax[cuda]&quot; -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
!pip install flax
!pip install optax

import jax

rng = jax.random.PRNGKey(0)"><pre class="notranslate">!p<span class="pl-s1">ip</span> <span class="pl-s1">install</span> <span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-s1">upgrade</span> <span class="pl-s1">pip</span>
!p<span class="pl-s1">ip</span> <span class="pl-s1">install</span> <span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-s1">upgrade</span> <span class="pl-s">"jax[cuda]"</span> <span class="pl-c1">-</span><span class="pl-s1">f</span> <span class="pl-s1">https</span>:<span class="pl-s1"></span><span class="pl-c1">//</span><span class="pl-s1">storage</span>.<span class="pl-s1">googleapis</span>.<span class="pl-s1">com</span><span class="pl-c1">/</span><span class="pl-s1">jax</span><span class="pl-c1">-</span><span class="pl-s1">releases</span><span class="pl-c1">/</span><span class="pl-s1">jax_cuda_releases</span>.<span class="pl-s1">html</span>
!p<span class="pl-s1">ip</span> <span class="pl-s1">install</span> <span class="pl-s1">flax</span>
!p<span class="pl-s1">ip</span> <span class="pl-s1">install</span> <span class="pl-s1">optax</span>

<span class="pl-k">import</span> <span class="pl-s1">jax</span>

<span class="pl-s1">rng</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-s1">random</span>.<span class="pl-v">PRNGKey</span>(<span class="pl-c1">0</span>)</pre></div>
<p dir="auto">and returns</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="---------------------------------------------------------------------------
XlaRuntimeError                           Traceback (most recent call last)
[&lt;ipython-input-2-71f91d108500&gt;](https://localhost:8080/#) in &lt;module&gt;
----&gt; 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`
-&gt; 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 "><pre class="notranslate"><code class="notranslate">---------------------------------------------------------------------------
XlaRuntimeError                           Traceback (most recent call last)
[&lt;ipython-input-2-71f91d108500&gt;](https://localhost:8080/#) in &lt;module&gt;
----&gt; 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`
-&gt; 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 
</code></pre></div>
<p dir="auto">Also reported here <a class="issue-link js-issue-link" data-error-text="Failed to load title" data-id="4856887" data-permission-text="Title is private" data-url="https://github.com/google/jax/discussions/14463" data-hovercard-type="discussion" data-hovercard-url="/google/jax/discussions/14463/hovercard" href="https://github.com/google/jax/discussions/14463#discussion-4856887">#14463 (comment)</a></p>
<h3 dir="auto">What jax/jaxlib version are you using?</h3>
<p dir="auto">jax == 0.4.3; jaxlib == 0.4.3+cuda11.cudnn86</p>
<h3 dir="auto">Which accelerator(s) are you using?</h3>
<p dir="auto">GPU</p>
<h3 dir="auto">Additional system info</h3>
<p dir="auto">Google colab</p>
<h3 dir="auto">NVIDIA GPU info</h3>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="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 |
+-----------------------------------------------------------------------------+"><pre class="notranslate"><code class="notranslate">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 |
+-----------------------------------------------------------------------------+
</code></pre></div>