<p dir="auto">I'm getting a weird error when trying to <code class="notranslate">jit</code> things:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="Traceback (most recent call last):
  File &quot;/home/ubuntu/dev/research/lottery/cifar10_convnet_run.py&quot;, line 228, in &lt;module&gt;
    test_make_batcher()
  File &quot;/home/ubuntu/dev/research/lottery/cifar10_convnet_run.py&quot;, line 87, in test_make_batcher
    for fn in [make_batcher(5, 2), jit(make_batcher(5, 2))]:
  File &quot;/home/ubuntu/dev/research/lottery/cifar10_convnet_run.py&quot;, line 79, in make_batcher
    splits = list(jnp.arange(1, num_examples // batch_size + 1) * batch_size)
  File &quot;/nix/store/ix119bzn54cgvsm98cidfdb4c8w0v1qc-python3.9-jax-0.3.0/lib/python3.9/site-packages/jax/_src/numpy/lax_numpy.py&quot;, line 6747, in deferring_binary_op
    return binary_op(self, other)
RuntimeError: UNKNOWN: PTX JIT compiler library not found
in external/org_tensorflow/tensorflow/stream_executor/cuda/cuda_asm_compiler.cc(56): 'cuLinkCreate(0, nullptr, nullptr, &amp;link_state)'"><pre class="notranslate"><code class="notranslate">Traceback (most recent call last):
  File "/home/ubuntu/dev/research/lottery/cifar10_convnet_run.py", line 228, in &lt;module&gt;
    test_make_batcher()
  File "/home/ubuntu/dev/research/lottery/cifar10_convnet_run.py", line 87, in test_make_batcher
    for fn in [make_batcher(5, 2), jit(make_batcher(5, 2))]:
  File "/home/ubuntu/dev/research/lottery/cifar10_convnet_run.py", line 79, in make_batcher
    splits = list(jnp.arange(1, num_examples // batch_size + 1) * batch_size)
  File "/nix/store/ix119bzn54cgvsm98cidfdb4c8w0v1qc-python3.9-jax-0.3.0/lib/python3.9/site-packages/jax/_src/numpy/lax_numpy.py", line 6747, in deferring_binary_op
    return binary_op(self, other)
RuntimeError: UNKNOWN: PTX JIT compiler library not found
in external/org_tensorflow/tensorflow/stream_executor/cuda/cuda_asm_compiler.cc(56): 'cuLinkCreate(0, nullptr, nullptr, &amp;link_state)'
</code></pre></div>
<p dir="auto">What does this mean? What .so am I missing? I can confirm that I have <code class="notranslate">libnvidia-ptxjitcompiler.so</code>, <code class="notranslate">libnvidia-ptxjitcompiler.so.1</code>, and <code class="notranslate">libnvidia-ptxjitcompiler.so.495.44</code> all on my system.</p>
<p dir="auto">Please:</p>
<ul class="contains-task-list">
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> Check for duplicate issues.</li>
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this:</li>
</ul>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="from jax import jit
jit(lambda x: 2*x)"><pre class="notranslate"><span class="pl-k">from</span> <span class="pl-s1">jax</span> <span class="pl-k">import</span> <span class="pl-s1">jit</span>
<span class="pl-en">jit</span>(<span class="pl-k">lambda</span> <span class="pl-s1">x</span>: <span class="pl-c1">2</span><span class="pl-c1">*</span><span class="pl-s1">x</span>)</pre></div>
<ul class="contains-task-list">
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> If applicable, include full error messages/tracebacks.<br>
<a href="https://gist.github.com/samuela/7245561c382842d70a729caecbd11912">https://gist.github.com/samuela/7245561c382842d70a729caecbd11912</a></li>
</ul>
<ul class="contains-task-list">
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> Check for duplicate issues.<br>
I searched <code class="notranslate">is:issue is:open pip install jaxlib metadata</code></li>
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> Provide a complete example of how to reproduce the bug:</li>
</ul>
<div class="highlight highlight-source-shell notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="pip install --upgrade jax==0.2.0 jaxlib==0.1.55+cuda110 -f https://storage.googleapis.com/jax-releases/jax_releases.html"><pre class="notranslate">pip install --upgrade jax==0.2.0 jaxlib==0.1.55+cuda110 -f https://storage.googleapis.com/jax-releases/jax_releases.html</pre></div>
<ul class="contains-task-list">
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> If applicable, include full error messages/tracebacks. Returns:</li>
</ul>
<blockquote>
<p dir="auto">Looking in links: <a href="https://storage.googleapis.com/jax-releases/jax_releases.html" rel="nofollow">https://storage.googleapis.com/jax-releases/jax_releases.html</a><br>
Requirement already satisfied: jax==0.2.0 in /home/jb/anaconda3/envs/cuda111/lib/python3.8/site-packages (0.2.0)<br>
Collecting jaxlib==0.1.55+cuda110<br>
Using cached <a href="https://storage.googleapis.com/jax-releases/cuda110/jaxlib-0.1.55%2Bcuda110-cp38-none-manylinux2010_x86_64.whl" rel="nofollow">https://storage.googleapis.com/jax-releases/cuda110/jaxlib-0.1.55%2Bcuda110-cp38-none-manylinux2010_x86_64.whl</a> (155.2 MB)<br>
WARNING: Discarding <a href="https://storage.googleapis.com/jax-releases/cuda110/jaxlib-0.1.55%2Bcuda110-cp38-none-manylinux2010_x86_64.whl" rel="nofollow">https://storage.googleapis.com/jax-releases/cuda110/jaxlib-0.1.55%2Bcuda110-cp38-none-manylinux2010_x86_64.whl</a> (from <a href="https://storage.googleapis.com/jax-releases/jax_releases.html" rel="nofollow">https://storage.googleapis.com/jax-releases/jax_releases.html</a>). Requested jaxlib==0.1.55+cuda110 from <a href="https://storage.googleapis.com/jax-releases/cuda110/jaxlib-0.1.55%2Bcuda110-cp38-none-manylinux2010_x86_64.whl" rel="nofollow">https://storage.googleapis.com/jax-releases/cuda110/jaxlib-0.1.55%2Bcuda110-cp38-none-manylinux2010_x86_64.whl</a> has inconsistent version: filename has '0.1.55+cuda110', but metadata has '0.1.55'<br>
ERROR: Could not find a version that satisfies the requirement jaxlib==0.1.55+cuda110 (from versions: 0.1, 0.1.1, 0.1.32, 0.1.33, 0.1.34, 0.1.35, 0.1.36, 0.1.37, 0.1.38, 0.1.39, 0.1.40, 0.1.41, 0.1.42, 0.1.43, 0.1.44, 0.1.45, 0.1.46, 0.1.47, 0.1.48, 0.1.50, 0.1.51, 0.1.52, 0.1.52+cuda100, 0.1.52+cuda101, 0.1.52+cuda102, 0.1.52+cuda110, 0.1.55, 0.1.55+cuda100, 0.1.55+cuda101, 0.1.55+cuda102, 0.1.55+cuda110, 0.1.56, 0.1.56+cuda100, 0.1.56+cuda101, 0.1.56+cuda102, 0.1.56+cuda110, 0.1.56+cuda111, 0.1.57, 0.1.57+cuda101, 0.1.57+cuda102, 0.1.57+cuda110, 0.1.57+cuda111, 0.1.58, 0.1.58+cuda101, 0.1.58+cuda102, 0.1.58+cuda110, 0.1.58+cuda111, 0.1.59, 0.1.59+cuda101, 0.1.59+cuda102, 0.1.59+cuda110, 0.1.59+cuda111, 0.1.60, 0.1.60+cuda101, 0.1.60+cuda102, 0.1.60+cuda110, 0.1.60+cuda111, 0.1.61, 0.1.61+cuda101, 0.1.61+cuda102, 0.1.61+cuda110, 0.1.61+cuda111, 0.1.61+cuda112, 0.1.62, 0.1.62+cuda101, 0.1.62+cuda102, 0.1.62+cuda110, 0.1.62+cuda111, 0.1.62+cuda112, 0.1.63, 0.1.63+cuda101, 0.1.63+cuda102, 0.1.63+cuda110, 0.1.63+cuda111, 0.1.63+cuda112, 0.1.64, 0.1.64+cuda101, 0.1.64+cuda102, 0.1.64+cuda110, 0.1.64+cuda111, 0.1.64+cuda112, 0.1.65, 0.1.65+cuda101, 0.1.65+cuda102, 0.1.65+cuda110, 0.1.65+cuda111, 0.1.65+cuda112, 0.1.66, 0.1.66+cuda101, 0.1.66+cuda102, 0.1.66+cuda110, 0.1.66+cuda111, 0.1.67, 0.1.67+cuda101, 0.1.67+cuda102, 0.1.67+cuda110, 0.1.67+cuda111)<br>
ERROR: No matching distribution found for jaxlib==0.1.55+cuda110</p>
</blockquote>
<p dir="auto">Conda env at time of install:<br>
<a href="https://github.com/google/jax/files/6626643/tmp.txt">env.txt</a></p>
<p dir="auto">Notable: my cuda toollkit version is 11.1 (to use with torch). I am pinned back on jax 0.1.55 due to DeepMind GLN implementation, if curious.</p>
<p dir="auto">Edit: Using GTX 1070</p>