Import statement has to differ across versions

Author: ae-foster

folx/experimental/pallas/attention/custom_gradients.py

@@ -5,7 +5,11 @@
 import jax
 import jax.numpy as jnp
 from jax.experimental import pallas as pl
-from jax.experimental.pallas import gpu as plgpu
+
+try:
+    from jax.experimental.pallas import triton as plgpu
+except ImportError:
+    from jax.experimental.pallas import gpu as plgpu
 
 from .mhsa import mhsa_kernel, reference_mhsa_kernel
 from .mhsea import mhsea_kernel, reference_mhsea_kernel

folx/experimental/pallas/attention/forward_laplacian.py

@@ -5,7 +5,11 @@
 import jax
 import jax.numpy as jnp
 from jax.experimental import pallas as pl
-from jax.experimental.pallas import gpu as plgpu
+
+try:
+    from jax.experimental.pallas import triton as plgpu
+except ImportError:
+    from jax.experimental.pallas import gpu as plgpu
 
 from folx import forward_laplacian
 from folx.api import FwdJacobian, FwdLaplArray

folx/experimental/pallas/attention/mhsa.py

@@ -5,7 +5,11 @@
 import jax
 import jax.numpy as jnp
 from jax.experimental import pallas as pl
-from jax.experimental.pallas import gpu as plgpu
+
+try:
+    from jax.experimental.pallas import triton as plgpu
+except ImportError:
+    from jax.experimental.pallas import gpu as plgpu
 
 from .utils import (
     big_number,

folx/experimental/pallas/attention/mhsea.py

@@ -5,7 +5,11 @@
 import jax
 import jax.numpy as jnp
 from jax.experimental import pallas as pl
-from jax.experimental.pallas import gpu as plgpu
+
+try:
+    from jax.experimental.pallas import triton as plgpu
+except ImportError:
+    from jax.experimental.pallas import gpu as plgpu
 
 from .utils import (
     big_number,