Add information about unpacked element from prims.unpack_sequence

[kshitij12345]

Fixes #1635

Example

import thunder
import torch

def fn(x):
    return x.sin()

jfn = thunder.jit(fn)
jfn(torch.randn(3, requires_grad=True))

backward_trc = thunder.last_backward_traces(jfn)[-1]
print(backward_trc)

Backward before PR

def backward_fn(saved_for_backward, cotangents):
  # saved_for_backward: "Collection"
  # cotangents: "Collection"
  C0, _, = saved_for_backward
  clear_mutable_collection(saved_for_backward)
  del saved_for_backward
  t2, = cotangents
  clear_mutable_collection(cotangents)
  del cotangents
  x, = C0
  clear_mutable_collection(C0)
  del C0
  bw_t6 = torch.cos(x)  # bw_t6: "cpu f32[3]"
    # bw_t6 = ltorch.cos(x)  # bw_t6: "cpu f32[3]"
      # bw_t6 = prims.cos(x)  # bw_t6: "cpu f32[3]"
  del x
  bw_t7 = torch.mul(t2, bw_t6)  # bw_t7: "cpu f32[3]"
    # bw_t7 = ltorch.mul(t2, bw_t6)  # bw_t7: "cpu f32[3]"
      # bw_t7 = prims.mul(t2, bw_t6)  # bw_t7: "cpu f32[3]"
  del t2, bw_t6
  return (bw_t7,)

Backward after PR

@torch.no_grad()
@no_autocast
def backward_fn(saved_for_backward, cotangents):
  # saved_for_backward: "Collection"
  # cotangents: "Collection"
  C0, _, = saved_for_backward
  # C0: "Collection"
  # None
  clear_mutable_collection(saved_for_backward)
  del saved_for_backward
  t2, = cotangents
  # t2: "cpu f32[3]"
  clear_mutable_collection(cotangents)
  del cotangents
  x, = C0
  # x: "cpu f32[3]"
  clear_mutable_collection(C0)
  del C0
  bw_t6 = torch.cos(x)  # bw_t6: "cpu f32[3]"
    # bw_t6 = ltorch.cos(x)  # bw_t6: "cpu f32[3]"
      # bw_t6 = prims.cos(x)  # bw_t6: "cpu f32[3]"
  del x
  bw_t7 = torch.mul(t2, bw_t6)  # bw_t7: "cpu f32[3]"
    # bw_t7 = ltorch.mul(t2, bw_t6)  # bw_t7: "cpu f32[3]"
      # bw_t7 = prims.mul(t2, bw_t6)  # bw_t7: "cpu f32[3]"
  del t2, bw_t6
  return (bw_t7,)

[t-vi]

Awesome, thank you @kshitij12345

[riccardofelluga]

Just a curiosity, why the info is displayed in the next row instead of on the side?

[t-vi]

Printing is a bit tricky for the sequence unpack, as it can be 100s of results.