From b3164230b638946190c1d8db2d34afdffd5681e2 Mon Sep 17 00:00:00 2001 From: Tai An Date: Sat, 5 Sep 2026 00:17:15 -0700 Subject: [PATCH] fix(jax): interpolate the tensor-sequence-parallelism warning The warning raised when tp_resource appears in lhs_specs is built from three implicitly concatenated fragments, none of which carries the f prefix, so users see the literal text {gsr.tp_resource} and {lhs_specs} instead of the values that triggered the warning. Both names are locals defined a few lines above. Signed-off-by: Tai An --- transformer_engine/jax/cpp_extensions/gemm.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/transformer_engine/jax/cpp_extensions/gemm.py b/transformer_engine/jax/cpp_extensions/gemm.py index 68c72d5059..3e2df50833 100644 --- a/transformer_engine/jax/cpp_extensions/gemm.py +++ b/transformer_engine/jax/cpp_extensions/gemm.py @@ -895,8 +895,8 @@ def _parse_operand_output_specs( if gsr.tp_resource is not None: if gsr.tp_resource in lhs_specs: warnings.warn( - "Tensor sequence parallelism is detected as tp_resource='{gsr.tp_resource}'" - " appears in lhs_specs: {lhs_specs}. Please setting MeshResource.tpsp_resource" + f"Tensor sequence parallelism is detected as tp_resource='{gsr.tp_resource}'" + f" appears in lhs_specs: {lhs_specs}. Please setting MeshResource.tpsp_resource" " for tensor sequence parallelism to avoid potential issues." )