Merge pull request #296 from dfm:gpu-support-error-message

PiperOrigin-RevId: 675257703
jax-ml · Sep 16, 2024 · 949f07b · 949f07b
2 parents 4ebaed0 + a5dc021
commit 949f07b
Showing 1 changed file with 10 additions and 1 deletion.
diff --git a/jax_triton/__init__.py b/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