Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Unable to Install #246

Open
karan-dalal opened this issue Jan 3, 2024 · 2 comments
Open

Unable to Install #246

karan-dalal opened this issue Jan 3, 2024 · 2 comments

Comments

@karan-dalal
Copy link

karan-dalal commented Jan 3, 2024

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
Copy link

Facing the same issue. Any updates on this?

@karan-dalal
Copy link
Author

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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

No branches or pull requests

2 participants