Constraints for jaxlib version during jax installation

In v0.2.11 jax added the nice optional installation targets cpu, cudaXXX and tpu which install jaxlib for the desired platform. Currently, the setup script guarantees that the installed version of jaxlib is compatible by constraining the version number from below to the earliest known compatible version. However, no similar constraint is put for a maximum compatible version number. Given that jaxlib does not appear to necessarily maintain backward compatibility between updates, this is an issue for installing older versions of jax (as these will always pull the latest available jaxlib version, leading to crashes due to incompatibility).

Jax installation should always result in a working setup, i.e., only the latest compatible version of jaxlib should be installed when installing jax. This either requires version constraints in jax’s setup.py or backward compatibility guarantees in jaxlib (or a combination of both).

Please:

• Check for duplicate requests.
• Describe your goal, and if possible provide a code snippet with a motivating example.

[lumip]

I don’t have the time right now, so please go ahead 😃

[jakevdp]

Sounds good - are you interested in sending a PR? If not I’m happy to take care of it.