[JAX] fix benchmarking function

cuequivariance_jax/cuequivariance_jax/benchmarking.py

@@ -19,7 +19,6 @@
 from functools import partial
 
 import jax
-from jax.experimental.mosaic.gpu.profiler import _event_elapsed, _event_record
 
 
 def measure_clock_ticks(f, *args, **kwargs) -> tuple[float, float]:
@@ -49,6 +48,18 @@ def my_function(x, y):
     """
     from cuequivariance_ops_jax import noop, sleep, synchronize
 
+    try:
+        from jax.experimental.mosaic.gpu.profiler import _event_elapsed, _event_record
+    except ImportError:
+        try:
+            from cuequivariance_ops_jax import event_elapsed as _event_elapsed
+            from cuequivariance_ops_jax import event_record as _event_record
+        except ImportError:
+            raise ImportError(
+                "_event_elapsed and/or _event_record not found in jax.experimental.mosaic.gpu.profiler\n"
+                "They are known to be available in jax>=0.4.36,<=0.7.1"
+            )
+
     def run_func(state):
         """Wrapper function that calls the target function and ensures proper data flow."""
         args, kwargs = state