Preserve tags when swapping proxies (#1189)

Author: ali-alshaar7

thunder/core/jit_ext.py

@@ -1577,6 +1577,11 @@ def _get_process_group_from(*fn_and_args) -> Optional["ProcessGroup"]:
     return found_pg
 
 
+def update_tags(proxy_swapmap: dict[Variable, Proxy]) -> None:
+    for old, new in proxy_swapmap.items():
+        new.tags.update(unvariableify(old).tags)
+
+
 def thunder_general_jit(
     fn: Callable,
     args: tuple[Any, ...],
@@ -1682,6 +1687,8 @@ def restrict_proxy_swapmap(proxies: tuple[Proxy]) -> dict[Variable, Proxy]:
     # Update prologue trace by renaming proxies which are passed from prologue to the computation trace
     prologue_trace = _apply_trace_proxy_rename(prologue_trace, restrict_proxy_swapmap(pro_to_comp_proxies))
 
+    update_tags(ctx._proxy_swapmap)
+
     # Update computation trace by renaming proxies which are in the ctx._proxy_swapmap
     computation_trace = _apply_trace_proxy_rename(computation_trace, ctx._proxy_swapmap, "computation")