<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"> The environment is ubuntu 20.04 &amp; Jetpack v5.0.1 - build follows Jetson Zoo website</li>
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox"> I have tried to both compile jax without success.  pip3 install jax[cuda] works.  However if I run pip3 install jaxlib I get an error code Error: No matching distribution found for jaxlib</li>
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox"> when I have tried to compile, the .cache for bazel uses all my free disk space.</li>
</ul>
<p dir="auto">Thanks in advance for your assistance,</p>
<p dir="auto">Paul</p>
<p dir="auto">Potentially related to <a class="issue-link js-issue-link" data-error-text="Failed to load title" data-id="916640295" data-permission-text="Title is private" data-url="https://github.com/google/jax/issues/6932" data-hovercard-type="issue" data-hovercard-url="/google/jax/issues/6932/hovercard" href="https://github.com/google/jax/issues/6932">#6932</a> or <a class="issue-link js-issue-link" data-error-text="Failed to load title" data-id="927516918" data-permission-text="Title is private" data-url="https://github.com/google/jax/issues/7052" data-hovercard-type="issue" data-hovercard-url="/google/jax/issues/7052/hovercard" href="https://github.com/google/jax/issues/7052">#7052</a>.</p>
<p dir="auto">I am using a heterogeneous cluster with various different compute options. I got Jax and Flax installed fine on the CPU system. Then I tried using a GPU node (specifically dual socket Cavium ThunderX2 99xx 32-core processors with two NVIDIA Tesla V100 GPUs), but am unable to install jaxlib:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="$ pip install --upgrade jax jaxlib==0.1.68+cuda111 -f https://storage.googleapis.com/jax-releases/jax_releases.html
Looking in links: https://storage.googleapis.com/jax-releases/jax_releases.html
Requirement already satisfied: jax in /home/southworth/.conda/envs/jax_arch/lib/python3.7/site-packages (0.2.16)
ERROR: Could not find a version that satisfies the requirement jaxlib==0.1.68+cuda111 (from versions: none)
ERROR: No matching distribution found for jaxlib==0.1.68+cuda111"><pre class="notranslate"><code class="notranslate">$ pip install --upgrade jax jaxlib==0.1.68+cuda111 -f https://storage.googleapis.com/jax-releases/jax_releases.html
Looking in links: https://storage.googleapis.com/jax-releases/jax_releases.html
Requirement already satisfied: jax in /home/southworth/.conda/envs/jax_arch/lib/python3.7/site-packages (0.2.16)
ERROR: Could not find a version that satisfies the requirement jaxlib==0.1.68+cuda111 (from versions: none)
ERROR: No matching distribution found for jaxlib==0.1.68+cuda111
</code></pre></div>
<p dir="auto">or even</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="pip install jaxlib
ERROR: Could not find a version that satisfies the requirement jaxlib (from versions: none)
ERROR: No matching distribution found for jaxlib"><pre class="notranslate"><code class="notranslate">pip install jaxlib
ERROR: Could not find a version that satisfies the requirement jaxlib (from versions: none)
ERROR: No matching distribution found for jaxlib
</code></pre></div>
<p dir="auto">Below is my pip version, which is up to date. Pip also installs other libraries fine (numpy, scipy, etc.)</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="$ pip --version
pip 21.1.2 from /home/southworth/.conda/envs/jax_arch/lib/python3.7/site-packages/pip (python 3.7)"><pre class="notranslate"><code class="notranslate">$ pip --version
pip 21.1.2 from /home/southworth/.conda/envs/jax_arch/lib/python3.7/site-packages/pip (python 3.7)
</code></pre></div>