search

jax-ml / jax-triton

jax-triton contains integrations between JAX and OpenAI Triton
Apache License 2.0
350 stars 40 forks source link

Stable set of dependencies? #284

Open vvvm23 opened 3 months ago

vvvm23 commented 3 months ago

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 commented 4 weeks ago

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