fix(jax): interpolate the tensor-sequence-parallelism warning

[Anai-Guo]

Problem

transformer_engine/jax/cpp_extensions/gemm.py, in _parse_operand_output_specs:

        gsr = global_mesh_resource()

        # Ensure that tensor sequence parallelism is not used via setting tp_resource
        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"
                    " for tensor sequence parallelism to avoid potential issues."
                )

The message is built from three implicitly concatenated fragments and none of them
carries the f prefix, so the user is shown the literal {gsr.tp_resource} and
{lhs_specs} — precisely the two values that would tell them which resource and which
spec caused the warning. Both are locals bound a few lines above.

Why this is a typo

gemm.py contains 78 string literals with {...} interpolation. 76 carry the f
prefix; the only two that do not are the two fragments of this one call.

Verification

I pulled the warnings.warn call out of both the current main file and the patched
file with ast (no hand-transcription) and evaluated each under
warnings.catch_warnings with gsr.tp_resource = "tp", lhs_specs = ("dp", "tp", None):

[BEFORE (main)]
  Tensor sequence parallelism is detected as tp_resource='{gsr.tp_resource}' appears in
  lhs_specs: {lhs_specs}. Please setting MeshResource.tpsp_resource for tensor sequence
  parallelism to avoid potential issues.

[AFTER  (patch)]
  Tensor sequence parallelism is detected as tp_resource='tp' appears in
  lhs_specs: ('dp', 'tp', None). Please setting MeshResource.tpsp_resource for tensor
  sequence parallelism to avoid potential issues.

The f prefix is added only to the two fragments that contain placeholders — the third
is left plain so it does not become an empty f-string.

Formatting checked with the pinned pre-commit hook exactly as configured
(black==24.4.2 --line-length=100 --preview --enable-unstable-feature=string_processing):
1 file left unchanged.

Left alone

The message also reads "Please setting MeshResource.tpsp_resource". I did not touch the
wording so this stays a single-concern change — happy to fix it if you'd like.

---

🤖 Generated with Claude Code

[greptile-apps]

Greptile Summary

This PR fixes a JAX tensor-sequence-parallelism warning so it displays the actual tensor-parallel resource and left-hand-side partition specification instead of literal placeholders.

• Adds f-string interpolation to the two warning fragments containing runtime values.
• Leaves warning conditions and GEMM behavior unchanged.

Confidence Score: 5/5

The PR appears safe to merge because it only corrects diagnostic interpolation without changing control flow or computation.

The interpolated values have safe built-in representations, and no actionable correctness, security, or quality issues remain.

Important Files Changed

Filename	Overview
transformer_engine/jax/cpp_extensions/gemm.py	Correctly interpolates the tensor-parallel resource and operand specification in an existing warning message.

_{Reviews (1): Last reviewed commit: "fix(jax): interpolate the tensor-sequenc..." | Re-trigger Greptile}