Unable to Install

[karan-dalal]

Hi there. I'm attempting to follow the installation directions. First, I install jax triton using pip install jax-triton.

Then, I install JAX and JAXLib using the Install JAX Docs:
pip install --upgrade "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

Currently, jax-triton installs JAX version 0.4.14. But there are no compatible JAXlib versions available. Thus, I install the most up to date JAX / JAXlib using the command above.

I keep getting the following error message when running the example code:
ValueError: triton_call is only available when triton is installed.

What is the correct / updated way to install JAX Triton? It would be great to have some clearer docs, I'm happy to make a PR once I can get it working.

[markschoene]

Facing the same issue. Any updates on this?

[karan-dalal]

Was unable to get this working. Would recommend trying out Pallas – google/jax#18603