pip installation: CPU on Linux workstation
JAX is written in pure Python, but it depends on XLA, which needs to be installed as the jaxlib package
python -m venv ~/opt/jax/jaxcpu
This command creates the target directory and any parent directoires that don't exist already and places a pyvenv.cfg file in it with a home key pointing ot the python installation from which the command was run.
pip install --upgrade pip
pip install --upgrade "jax[cpu]"
pip installation: GPU on Linux workstation