[JAX][Fix] Gate collective-stream compute_on on import and silence pylint

transformer_engine/jax/cpp_extensions/ep.py

@@ -42,7 +42,7 @@ def wrapper(*args, **kwargs):
         # a nullary thunk that closes over them; the array operands are captured
         # as consts and lifted to real operands, outputs stay on device. XLA
         # async-wraps the resulting call onto the collective stream.
-        annotated = compute_on(
+        annotated = compute_on(  # pylint: disable=not-callable
             compute_type="gpu_stream:collective",
             out_memory_spaces=jax.memory.Space.Device,
         )(lambda: func(*args, **kwargs))

transformer_engine/jax/version_utils.py

@@ -69,9 +69,16 @@ def is_triton_autotuned_alias_safe() -> bool:
 _COLLECTIVE_STREAM_MIN_JAX_VERSION = "0.10.1"
 
 
+@lru_cache(maxsize=None)
 def is_collective_stream_supported() -> bool:
     """Return True if the installed JAX supports the gpu_stream:collective annotation."""
-    return jax_version_meet_requirement(_COLLECTIVE_STREAM_MIN_JAX_VERSION)
+    if not jax_version_meet_requirement(_COLLECTIVE_STREAM_MIN_JAX_VERSION):
+        return False
+    try:
+        from jax.experimental.compute_on import compute_on  # pylint: disable=unused-import
+    except ImportError:
+        return False
+    return True
 
 
 def is_triton_extension_supported() -> bool: