Install jaxlib from a Nextstrain-managed repository

[victorlin]

Originally from #50 (comment)

On the linux/arm64 platform, jaxlib is installed from a pre-built binary provided by an external GitHub user:

docker-base/Dockerfile

Lines 172 to 179 in 960fb5b

	# Install jaxlib on linux/arm64
	# jaxlib, an evofr dependency, does not have official pre-built binaries for
	# linux/arm64. A GitHub user has provided them in a fork repo.
	# https://github.com/google/jax/issues/7097#issuecomment-1110730040
	RUN if [[ "$TARGETPLATFORM" == linux/arm64 ]]; then \
	pip3 install https://github.com/yoziru/jax/releases/download/jaxlib-v0.3.25/jaxlib-0.3.25-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.whl \
	; \
	fi

We should look into building and repackaging this library ourselves for security and reliability reasons.

[victorlin]

An example of how the unofficial binary is not reliable: #173 (comment)

[victorlin]

I'm looking into building JAX locally.

[victorlin]

I have an unfinished attempt at victorlin/build-jax. Summarizing the progress here since I plan to focus on other things for now.

The goal was to try to compile an AARCH64 version of JAX on a GitHub-hosted Linux runner (AMD64).

I briefly tried building via Docker under emulation, which was unsuccessful. Even if it were successful, it would probably take longer than the maximum allowed GitHub Actions workflow run time.

Then, I focused on cross-compiling since that has been done before. I attempted this first with then without Docker (for easier debugging of individual commands as GitHub Actions steps). With a 3.5 hour run time, I was able to build a file jaxlib-0.4.9-cp310-cp310-manylinux2014_aarch64.whl (a downloadable artifact of the linked run). However, I'm not sure that this wheel works because auditwheel repair failed on the file. I think the failure makes sense because I didn't use a manylinux Docker image. To address that, I tried building with messense/manylinux2014-cross:x86_64 but had troubles running the setup commands before the actual JAX build.

[victorlin]

I just noticed that conda-forge includes an AARCH64 version of jaxlib (feedstock). It's not maintained by JAX developers, but at least looks more reliable than the status quo.

We might not want to use it directly due to reasons described in #127 (comment), but it's one example that would benefit from using Conda in the image build.

[victorlin]

This is no longer necessary now that jaxlib version 0.4.18 includes Linux aarch64 wheels. I've pushed 85fe9a7 to try things out and will make a PR if it goes well.