[Bug] [CI] Error installing Jax for Ubuntu ci_cpu Docker build

[Liam-Sturge]

Expected behavior

When building an Ubuntu cpu container the script docker/install/ubuntu_install_jax.sh should install Jaxlib correctly without issues.

Actual behavior

Jaxlib fails to install with the following error reported:

Collecting orbax
  Using cached orbax-0.1.8.tar.gz (1.6 kB)
    ERROR: Command errored out with exit status 1:
     command: /usr/bin/python3 -c 'import sys, setuptools, tokenize; sys.argv[0] = '"'"'/tmp/pip-install-sy4y9jni/orbax/setup.py'"'"'; __file__='"'"'/tmp/pip-install-sy4y9jni/orbax/setup.py'"'"';f=getattr(tokenize, '"'"'open'"'"', open)(__file__);code=f.read().replace('"'"'\r\n'"'"', '"'"'\n'"'"');f.close();exec(compile(code, __file__, '"'"'exec'"'"'))' egg_info --egg-base /tmp/pip-install-sy4y9jni/orbax/pip-egg-info
         cwd: /tmp/pip-install-sy4y9jni/orbax/
    Complete output (3 lines):
    
    *** Orbax is a namespace, and not a standalone package. For model checkpointing and exporting utilities, please install `orbax-checkpoint` and `orbax-export` respectively (instead of `orbax`). ***
    
    ----------------------------------------
ERROR: Command errored out with exit status 1: python setup.py egg_info Check the logs for full command output.

This is due to the latest release of Orbax (v0.1.8) which enforces the deprecation of installing orbax as a standalone package. The version of Flax that docker/install/ubuntu_install_jax.sh attempts to install (v0.6.8) has orbax rather than orbax-checkpoint specified as an installation dependency.

Environment

An x86 platform running Ubuntu 22.04.

Steps to reproduce

Occurs when setting up an Ubuntu 22.04 cpu container using Dockerfile.ci_cpu.

Can also be reproduced by running the following commands in an Ubuntu environment:

pip3 install jaxlib==0.4.7 "jax[cpu]==0.4.7"
pip3 install flax==0.6.8

Triage

• dev:ci

cc @Mousius @areusch @leandron

[Liam-Sturge]

No longer seeing the Installation error after bumping the version of Jaxlib and Flax.