Provide AArch64 (ARM) Linux jaxlib wheels

Potentially related to https://github.com/google/jax/issues/6932 or https://github.com/google/jax/issues/7052.

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:

$ 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

or even

pip install jaxlib
ERROR: Could not find a version that satisfies the requirement jaxlib (from versions: none)
ERROR: No matching distribution found for jaxlib

Below is my pip version, which is up to date. Pip also installs other libraries fine (numpy, scipy, etc.)

$ pip --version
pip 21.1.2 from /home/southworth/.conda/envs/jax_arch/lib/python3.7/site-packages/pip (python 3.7)

[yoziru]

☝️ this indicates you’re running out of memory, especially since the 0.3.x release, you need pretty much the entire 16GB available.

I managed to build manylinux aarch64 wheels on M1, you can get them here: https://github.com/yoziru/jax/releases/tag/jaxlib-v0.3.5

Here is the Dockerfile for reproducing these wheels (run from the root of this repo after cloning):

FROM quay.io/pypa/manylinux2014_aarch64

COPY .bazelversion .
RUN BAZEL_VERSION=$(cat .bazelversion) \
    && curl -SL https://github.com/bazelbuild/bazel/releases/download/$BAZEL_VERSION/bazel-$BAZEL_VERSION-linux-arm64 \
    -o /usr/local/bin/bazel \
    && chmod +x /usr/local/bin/bazel

ARG PYTHON_VERSION_MAIN=38 \
    JAXLIB_VERSION=0.3.5

RUN ln -s /opt/python/cp${PYTHON_VERSION_MAIN}-cp${PYTHON_VERSION_MAIN}/bin/pip /usr/local/bin/pip \
    && ln -s /opt/python/cp${PYTHON_VERSION_MAIN}-cp${PYTHON_VERSION_MAIN}/bin/python /usr/local/bin/python \
    && ln -s /opt/python/cp${PYTHON_VERSION_MAIN}-cp${PYTHON_VERSION_MAIN}/bin/python /usr/local/bin/python3
RUN pip install numpy six
WORKDIR /builder
COPY . .
RUN python build/build.py
RUN auditwheel repair dist/jaxlib-${JAXLIB_VERSION}-cp${PYTHON_VERSION_MAIN}-none-manylinux2014_aarch64.whl

[markjens]

I am glad to see there is some progress already on this task! I hope we get aarch64 wheels soon at https://storage.googleapis.com/jax-releases/jax_cuda_releases.html !