How to install jax and jaxlib by using github url

[alcarazolabs]

I have python 3.9 in my device. When I run:

pip install jax and jaxlib
This is installing the version jax-0.4.21 jaxlib-0.4.21 However I need the version 0.4.5 for both jax and jaxlib.

When I run:

pip install jax==0.4.5 jaxlib==0.4.5

This is throwing this error:

ERROR: Could not find a version that satisfies the requirement jaxlib==0.4.5 (from versions: 0.4.18, 0.4.19, 0.4.20, 0.4.21)
    ERROR: No matching distribution found for jaxlib==0.4.5

I don't kwno why happens it. Lately I decide to install jax by using the git url, of this way:

pip install "git+https://github.com/google/[email protected]"

This throwed this message:

Collecting git+https://github.com/google/[email protected]
  Cloning https://github.com/google/jax.git (to revision jax-v0.4.5) to /tmp/pip-req-build-828jle88
  Running command git clone --filter=blob:none --quiet https://github.com/google/jax.git /tmp/pip-req-build-828jle88
  Running command git checkout -q cafaa50b25515a554568db06667b0c9b6abaff27
  Resolved https://github.com/google/jax.git to commit cafaa50b25515a554568db06667b0c9b6abaff27
  Preparing metadata (setup.py) ... done
Requirement already satisfied: numpy>=1.20 in ./.myenv/lib/python3.9/site-packages (from jax==0.4.5) (1.26.2)
Requirement already satisfied: opt_einsum in ./.myenv/lib/python3.9/site-packages (from jax==0.4.5) (3.3.0)
Requirement already satisfied: scipy>=1.5 in ./.myenv/lib/python3.9/site-packages (from jax==0.4.5) (1.11.4)
Building wheels for collected packages: jax
  Building wheel for jax (setup.py) ... done
  Created wheel for jax: filename=jax-0.4.5-py3-none-any.whl size=1424415 sha256=aff60a0ed1dea1d004cb5748ed322b9db90cd0064e470567f102a2f83d1873d1
  Stored in directory: /tmp/pip-ephem-wheel-cache-e8920quj/wheels/47/59/88/cceb9c59d0d692b940160f055bae0c60cd1295e4edc393ff48
Successfully built jax
Installing collected packages: jax
ERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.
chex 0.1.85 requires jaxlib>=0.1.37, which is not installed.
optax 0.1.7 requires jaxlib>=0.1.37, which is not installed.
orbax-checkpoint 0.4.6 requires jaxlib, which is not installed.
chex 0.1.85 requires jax>=0.4.16, but you have jax 0.4.5 which is incompatible.
flax 0.7.5 requires jax>=0.4.19, but you have jax 0.4.5 which is incompatible.
orbax-checkpoint 0.4.6 requires jax>=0.4.9, but you have jax 0.4.5 which is incompatible.

This seems that first we need to install jaxlib, but I don't know how install this through github url. I just know that (https://github.com/google/jax/releases/tag/jax-v0.4.5) is the github repo for jax but I don't know where is the github repo for jaxlib.

I will appreciate any idea to fix this problem guys, thanks so much.

[hawkinsp]

jaxlib is inside the JAX repository, but you can't build it with pip. If you don't use a binary package, you'd have to build from source (https://jax.readthedocs.io/en/latest/developer.html). I recommend using a binary package!

I don't think we ever actually released a jaxlib 0.4.5. We don't always release jaxlib when we release jax. I think you should use jaxlib 0.4.4 instead in this case: the most recent preceding jaxlib will be supported for any given jax.

Hope that helps!

[jakevdp]

+1 – also this page is useful if you're lookiing for older jaxlibs: https://pypi.org/project/jaxlib/#history. For info on installing versions older than those available on PyPI, see https://jax.readthedocs.io/en/latest/installation.html#installing-older-jaxlib-wheels

[alcarazolabs]

jaxlib is inside the JAX repository, but you can't build it with pip. If you don't use a binary package, you'd have to build from source (https://jax.readthedocs.io/en/latest/developer.html). I recommend using a binary package!

I don't think we ever actually released a jaxlib 0.4.5. We don't always release jaxlib when we release jax. I think you should use jaxlib 0.4.4 instead in this case: the most recent preceding jaxlib will be supported for any given jax.

Hope that helps!

Ok. I decided build from source but I got one error.

1. First I downloaded this source code from here: https://github.com/google/jax/releases/tag/jax-v0.4.5 I just downloaded the Source code
   (zip).

2. I unzipped it.

3. The I went inside the directory and I ran:

python build/build.py
And after 1 hour I got:

Bazel binary path: ./bazel-5.1.1-linux-arm64
Bazel version: 5.1.1
Python binary path: /home/orin/Documents/whisper/whisperjax/.myenv/bin/python
Python version: 3.9
NumPy version: 1.26.2
MKL-DNN enabled: yes
Target CPU: aarch64
Target CPU features: release
CUDA enabled: no
TPU enabled: no
Remote TPU enabled: no
ROCm enabled: no
Plugin device enabled: no
....
....
...
# Configuration: 92012d0c1249a20cab353cfe390d347cf69b0c5d540dca34fb7af7c76c995555
# Execution platform: @local_execution_config_platform//:platform
gcc: fatal error: Killed signal terminated program cc1plus
compilation terminated.
Target //build:build_wheel failed to build
INFO: Elapsed time: 858.305s, Critical Path: 108.11s
INFO: 569 processes: 13 internal, 556 local.
FAILED: Build did NOT complete successfully
ERROR: Build failed. Not running target
FAILED: Build did NOT complete successfully
b''
Traceback (most recent call last):
  File "/home/orin/Documents/whisper/whisperjax/jax045/jax-jax-v0.4.5/build/build.py", line 567, in <module>
    main()
  File "/home/orin/Documents/whisper/whisperjax/jax045/jax-jax-v0.4.5/build/build.py", line 562, in main
    shell(command)
  File "/home/orin/Documents/whisper/whisperjax/jax045/jax-jax-v0.4.5/build/build.py", line 53, in shell
    output = subprocess.check_output(cmd)
  File "/usr/lib/python3.9/subprocess.py", line 424, in check_output
    return run(*popenargs, stdout=PIPE, timeout=timeout, check=True,
  File "/usr/lib/python3.9/subprocess.py", line 528, in run
    raise CalledProcessError(retcode, process.args,
subprocess.CalledProcessError: Command '['./bazel-5.1.1-linux-arm64', 'run', '--verbose_failures=true', ':build_wheel', '--', '--output_path=/home/orin/Documents/whisper/whisperjax/jax045/jax-jax-v0.4.5/dist', '--cpu=aarch64']' returned non-zero exit status 1.
(.myenv) orin@ubuntu:~/Documents/whisper/whisperjax/jax045/jax-jax-v0.4.5$ 

I'm running it inside a nvidia jetson orin nano which is aarch64 architecture. The sistem has already installed bazel-1.2.1 which is compatible for aarch64.

Do you have any idea about what is happening? I will appreciate it. Thanks in advance.

PD: I also tried to run: pip install jax==0.4.5 jaxlib==0.4.4 However I'm getting the same error: see in this image: imgur.com/a/4Z3ITZv
It throwed :

ERROR: Could not find a version that satisfies the requirement jaxlib==0.4.4 (from versions: 0.4.18, 0.4.19, 0.4.20, 0.4.21) ERROR: No matching distribution found for jaxlib==0.4.4

[jakevdp]

If you want to build jaxlib version 0.4.4, you should use the jaxlib-v0.4.4 tag.

[jakevdp]

Regarding the error when installed from PyPI: what Python version are you using? jaxlib 0.4.4 only has pre-built wheels for CPython versions 3.8-3.11; if you're using Python 3.12 you won't be able to install jaxlib older than 0.4.18 or so.

[hawkinsp]

We also only started shipping aarch64 wheels as of jax 0.4.18: https://github.com/google/jax/blob/main/CHANGELOG.md#jaxlib-0418-oct-6-2023

[hawkinsp]

Although I'm not 100% sure if that aarch64 is compatible with jetson or not.

[alcarazolabs]

I fixed it running;
$ pip install jax jaxlib
$ pip install --upgrade jax jaxlib flax