[Pallas TPU] `jax.named_scope` splits otherwise co-issued instructions into separate bundles

Description

Minimal reproducer:

from contextlib import nullcontext

import jax
import jax.numpy as jnp
from jax.experimental import pallas as pl

def make_kernel(named):
    def kernel(x_ref, add_ref, mul_ref):
        x = x_ref[...]
        with jax.named_scope("add") if named else nullcontext():
            added = x + 1.0
        with jax.named_scope("mul") if named else nullcontext():
            scaled = x * 2.0
        add_ref[...] = added
        mul_ref[...] = scaled

    return pl.pallas_call(
        kernel,
        out_shape=(jax.ShapeDtypeStruct((8, 128), jnp.float32),) * 2,
        name=f"scope_{'named' if named else 'plain'}",
    )

x = jax.ShapeDtypeStruct((8, 128), jnp.float32)
for named in (False, True):
    jax.jit(make_kernel(named)).lower(x).compile()

Inspect the scope_plain and scope_named final bundle files. Relevant instructions, with unrelated instructions and operands omitted:

scope_plain:
0x12: { vadd.f32 ... ;; vmul.f32 ... }

scope_named:
0x12: { vtrace ... }  // start add
0x13: { vadd.f32 ... }
0x14: { vtrace ... }  // end add
0x15: { vtrace ... }  // start mul
0x16: { vmul.f32 ... }
0x17: { vtrace ... }  // end mul

Without scopes, vadd and vmul share one bundle. With scopes, additional vtrace instructions separate them.

Expected behavior

Adding naming annotations should preserve the ability to issue these operations in the same bundle.

Actual behaviour

Scope tracing inserts execution boundaries that separate the operations.

Explanation

jax.named_scope should attach source-name metadata to operations, allowing profiling tools to attribute compiled instructions or execution events to a user-defined scope. In XPlane, this attribution can be represented through metadata such as tf_op on existing events. It does not require a separate timed region for every scope. A scope may span multiple bundles, and instructions from different scopes may share a bundle.

Currently, scope boundaries are lowered into vtrace instructions that act as scheduling barriers. This imposes an execution-order constraint on what should merely be a naming annotation and alters the program being observed, which is therefore incorrect.

This was originally reported in #40637, but that issue was closed for reasons unclear to me, so I’m opening a new issue to follow up.

System info (python version, jaxlib version, accelerator, etc.)

jax:    0.11.2.dev20260913+adb0562417
jaxlib: 0.11.1
numpy:  2.5.3
python: 3.14.7 free-threading build (main, Sep  1 2026, 14:17:47) [Clang 22.1.3 ]
device info: TPU v4-4, 4 local devices"
process_count: 1
platform: uname_result(system='Linux', node='t1v-n-c577c7ce-w-0', release='5.19.0-1022-gcp', version='#24~22.04.1-Ubuntu SMP Sun Apr 23 09:51:08 UTC 2023', machine='x86_64')
libtpu: 0.0.48.dev20260912-nightly