ability to load jax & jaxlib on a Nvidia Xavier running Jetpack v5.0.1

[pdavis39]

Please:

• The environment is ubuntu 20.04 & Jetpack v5.0.1 - build follows Jetson Zoo website
• 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
• when I have tried to compile, the .cache for bazel uses all my free disk space.

Thanks in advance for your assistance,

Paul

[hawkinsp]

We currently only provide x86_64 wheels, which is why the pip installation does not work. There aren't a lot of alternatives to building from source at the moment. You could try building from source on a different aarch64 machine with more disk space, or cross-compile.

Closing as a duplicate of #7097