Improve the error message when jaxlib is installed without GPU support.

jax-ml / jax-triton

jax-triton contains integrations between JAX and OpenAI Triton

Apache License 2.0

354 stars 40 forks source link

Improve the error message when jaxlib is installed without GPU support. #296

Closed dfm closed 2 months ago

dfm commented 2 months ago

Currently jax-triton fails to import with an AttributeError when jaxlib isn't installed with CUDA or ROCM support. This error message should provide more useful feedback.