Stable set of dependencies?

[vvvm23]

Hi,

I feel I am not alone in rangling dependencies for jax-triton. Could the authors share a set of fixed versions of jax, jaxlib, triton and jax-triton that are known to be working with python 3.10?

One cannot use pip install jax-triton as the pypi package does not support python 3.10. So we are forced to install from HEAD in this case.

However, what "HEAD" is varies greatly, and there is a high probability at any one time that installing from HEAD for all these dependencies will lead to some failure.

Thanks

[sbhavani]

https://github.com/NVIDIA/JAX-Toolbox includes a nightly container for jax-triton (ghcr.io/nvidia/jax:triton) with open Dockerfiles. That could be a good reference for a stable set of dependencies.

I see the following versions in the latest nightly:

Python version: 3.10.12
jax-triton version: 0.2.0
triton version: 3.1.0
jax version: 0.4.36.dev20241103+2d94914c8