Create Arm wheels for Raspberry Pi

[bennuttall]

I maintain a package repository where we build Arm wheels of all the packages on PyPI for Raspberry Pi users: https://www.piwheels.org

We recently had an issue reported that the jax package is missing from the repo piwheels/packages#44 and I want to see if we can fix that.

We build from source distribution files when it's available on PyPI, but we're obviously not able to automate building projects like this with a complex build procedure who don't support building from source using the standard tooling. We build opencv ourselves and import wheels for it, and the tensorflow team build Arm wheels and send them to me for inclusion in the repo.

We've had 161k downloads of Tensorflow, it's one of our most popular packages - and without the inclusion of the wheel, Raspberry Pi users wouldn't be able to use it. https://www.piwheels.org/project/tensorflow/

Would you be interested in officially supporting Raspberry Pi? If so, all you need to do is build a wheel targeting the Arm platform, and send it to me as part of your release process.

If you have any questions, let me know and I'll do what I can to help.

Thanks

[hawkinsp]

Jax on the Raspberry Pi sounds fun!

I suspect the best option is to cross-compile jaxlib for Raspberry Pi; the build probably wouldn't work very well on Raspberry-Pi class hardware. Most likely we can adapt the approach TensorFlow takes: https://www.tensorflow.org/install/source_rpi

Getting jaxlib building in a cross-compilation setting is in the "contributions welcome" category; we're only a very small team so we most likely won't be able to get to this any time soon.

[MatthewGerber]

@hawkinsp I'd like to use JAX on my Raspberry Pi; however, I'm not too familiar with cross-compilation. In your estimation, how difficult would it be to cross-compile JAX/jaxlib for the 32-bit ARM CPU that is used in the Raspberry Pi 4? I looked at the TF pages on cross-compilation for Raspberry Pi, but it's not clear to me how to proceed with JAX/jaxlib. I'm happy to work on this given a few pointers. I think JAX would be a nice addition to the Pi ecosystem, as @bennuttall says. Thoughts?

[hawkinsp]

We're not going to add support for 32-bit architectures on any platform, I strongly suspect.

However, these days, we ship aarch64 Linux wheels, which can be installed on newer Raspberry Pi models.

[boredcoder411]

I'm not sure these wheels actually work, I saw on the pypi file listing that there were aarch64 wheels, but pip install jaxlib on my raspberry pi returns an error