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

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.

jax_triton/__init__.py

@@ -34,7 +34,16 @@
 from jax_triton.version import __version__
 from jax_triton.version import __version_info__
 
-get_compute_capability = gpu_triton.get_compute_capability
+try:
+ get_compute_capability = gpu_triton.get_compute_capability
+except AttributeError:
+ raise ImportError(
+ "jax-triton requires JAX to be installed with GPU support. The "
+ "installation page on the JAX documentation website includes "
+ "instructions for installing a supported version:\n"
+ "https://jax.readthedocs.io/en/latest/installation.html"
+ )
+
 if jaxlib.version.__version_info__ >= (0, 4, 14):
  try:
  get_serialized_metadata = gpu_triton.get_serialized_metadata