tensorflow[and-cuda] 2.15.0/2.15.1 compatibility with jax[cuda12]

[attaluris]

Issue type

Bug

Have you reproduced the bug with TensorFlow Nightly?

No

Source

source

TensorFlow version

2.15.0/2.15.1

Custom code

Yes

OS platform and distribution

Debian Bulleye

Mobile device

No response

Python version

3.9/3.10

Bazel version

No response

GCC/compiler version

No response

CUDA/cuDNN version

12.2

GPU model and memory

v100

Current behavior?

Hey y'all! I think tensorflow[and-cuda] is incompatible with jax[cuda12] and just wanted to clarify if this was expected.
The solve error I'm getting is:

Because tensorflow[and-cuda] (2.15.0) depends on nvidia-nccl-cu12 (2.16.5)
 and jax[cuda12] (0.4.23) depends on nvidia-nccl-cu12 (>=2.18.3), tensorflow[and-cuda] (2.15.0) is incompatible with jax[cuda12] (0.4.23).
So, because hex-packages depends on both jax[cuda12] (0.4.23) and tensorflow[and-cuda] (2.15.0), version solving failed.

and none of the jax[cuda12] versions with GPU compatibility support nvidia-nccl-cu12=2.16.5; does this requirement need to be hard or can it be looser to accomodate higher versions of nvidia-nccl-cu12?

Standalone code to reproduce the issue

[tool.poetry]
name = "test-jax-and-poetry"
version = "0.1.0"
description = ""
authors = ["Tim Nonet <tnonet@hex.tech>"]
readme = "README.md"

[tool.poetry.dependencies]
python = ">=3.9,<3.11"
tensorflow = { "version" = "2.15.1", extras = ["and-cuda"] }
jax = {"version" = "*", extras = ["cuda12"] }


[build-system]
requires = ["poetry-core"]
build-backend = "poetry.core.masonry.api"

Relevant log output

Because tensorflow[and-cuda] (2.15.0) depends on nvidia-nccl-cu12 (2.16.5)
 and jax[cuda12] (0.4.23) depends on nvidia-nccl-cu12 (>=2.18.3), tensorflow[and-cuda] (2.15.0) is incompatible with jax[cuda12] (0.4.23).
So, because hex-packages depends on both jax[cuda12] (0.4.23) and tensorflow[and-cuda] (2.15.0), version solving failed.

whenever I use a version of jax that has the cuda12 extra

[sushreebarsa]

@attaluris TensorFlow[and-cuda] 2.15.0/2.15.1 is likely not compatible with jax[cuda12]. There's a version mismatch with respect to the NVIDIA NCCL library, a component needed for GPU support in both TensorFlow and JAX.
TensorFlow 2.15.0/2.15.1 might depend on an older NCCL version (e.g., nvidia-nccl-cu12 version 2.16.5).
For any further queries please raise an issue in Jax repository.

Thank you!

[attaluris]

@sushreebarsa Is there a way we could loosen the strict requirement in tensorflow from =2.16.5 to >2.16.5?

I also raised an issue in the Jax repo

[sushreebarsa]

@attaluris Could you try to update tensorflow to 2.16.1 and jax to 0.4.28?
TF 2.16.1 uses nvidia-nccl-cu12==2.19.3 (from tensorflow[and-cuda]) which is compatible with JAX's requirement of >=2.18.3.

Thank you!

[attaluris]

@sushreebarsa thanks for the detail! 👀
I'm seeing the same issue as #63362 with tensorflow 2.16.1 so I'm using 2.15.1
can I get a confirmation that the tensorflow 2.16.1 installation will be fixed and maybe a timeline? thanks in advance!

[sushreebarsa]

@attaluris On ubuntu 20.04LTS machine with Tesla P100 GPU, we tried to install JAX 0.4.28 cuda version first and then installed tensorflow 2.15.1. Both were installed successfully. JAX and tf were able to detect GPUs as well. Please have a look at the below screenshot for reference.

Thank you!

[github-actions]

This issue is stale because it has been open for 7 days with no activity. It will be closed if no further activity occurs. Thank you.

[github-actions]

This issue was closed because it has been inactive for 7 days since being marked as stale. Please reopen if you'd like to work on this further.

[google-ml-butler]

Are you satisfied with the resolution of your issue?
Yes
No