From 24629c645ee8bbbbf5d2247af841d6bf1d9ac2d4 Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Fri, 31 Jul 2026 09:46:19 -0700 Subject: [PATCH] Bump collective stream min JAX version to 0.10.1 Signed-off-by: Phuong Nguyen --- transformer_engine/jax/version_utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/transformer_engine/jax/version_utils.py b/transformer_engine/jax/version_utils.py index 8765a4fce0..e797e83440 100644 --- a/transformer_engine/jax/version_utils.py +++ b/transformer_engine/jax/version_utils.py @@ -65,8 +65,8 @@ def is_triton_autotuned_alias_safe() -> bool: # XLA gained the ``gpu_stream:collective`` stream annotation in openxla/xla#39604, -# which ships in the JAX 0.10.0 release. Older XLA fatally fails on it. -_COLLECTIVE_STREAM_MIN_JAX_VERSION = "0.10.0" +# first shipping in JAX 0.10.1. Older XLA fatally fails on it. +_COLLECTIVE_STREAM_MIN_JAX_VERSION = "0.10.1" def is_collective_stream_supported() -> bool: