Do not propagate name_stacks into lower_jaxpr_to_fun.
Currently name_stacks have a somewhat confusing semantics, that do not work correctly in the presence of non-inlined functions.
This change clarifies the semantics of name_stacks:
* name_stacks as assigned to jaxpr equations during trace time are always rooted at the enclosing jaxpr. Names do not flow from the caller of a jaxpr into the jaxpr. Reason: in general a jaxpr might have multiple callers, and any context from one caller need not apply to other call sites.
* when inlining a jaxpr into another jaxpr, we should concatenate the name stacks.
* in places during HLO lowering where we in effect "inline" a subjaxpr into an enclosing function, we should propagate name stack information from callers to callees.
* when we emit an out-of-line function during HLO lowering, we must *not* propagate name stacks. (This is the key change: previously we did propagate name_stacks here, and we were wrong to do so.)
Separately, we have taught XLA's inliner to concatenate operator names with a `/` when inlining, which means that if XLA inlines functions later, it will have the same semantics as the "inline" subjaxpr behavior above.
To enforce this, we remove the name_stack argument from mlir.lower_jaxpr_to_fun(): that function emits an out-of-line HLO function, and hence it is never correct to propagate a name_stack in that case. We add a small amount of special handling for the main_function to inject a non-empty name stack, preserving the top-level `jit(f)` users may expect to see today.
In a future change, we should give tracebacks attached to jaxpr equations the same semantics.
PiperOrigin-RevId: 777017036
