From 18948c98aa047d05fb8509ac9c4e0fc973d4c3ef Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Fri, 7 Aug 2026 13:57:11 -0700 Subject: [PATCH] [JAX] Gate collective-stream compute_on on import and silence pylint not-callable Signed-off-by: Phuong Nguyen --- transformer_engine/jax/cpp_extensions/ep.py | 2 +- transformer_engine/jax/version_utils.py | 9 ++++++++- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/transformer_engine/jax/cpp_extensions/ep.py b/transformer_engine/jax/cpp_extensions/ep.py index 0b18c3bbb6..ca70ea145c 100644 --- a/transformer_engine/jax/cpp_extensions/ep.py +++ b/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)) diff --git a/transformer_engine/jax/version_utils.py b/transformer_engine/jax/version_utils.py index e797e83440..9d73ecd3b5 100644 --- a/transformer_engine/jax/version_utils.py +++ b/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: