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

[attaluris]

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

[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