Unhelpful variable renaming

[shino16]

The commit 7c916c13675bb05b1a5522a9c797b33e997e4f19 tries to give identifiers observed in the source to the proxies of variables. However such identifiers sometimes refer to thunder's internal implementation of the language's construct, and this can be confusing.

Code sample

@thunder.jit
def f(xs, s):
    for i, x in enumerate(xs):
        s += x
    return s

n = 6
xs = [torch.zeros(n) for _ in range(n)]
s = torch.zeros(n)
f(xs, s)
print(thunder.last_traces(f)[-1])

@torch.no_grad()
@no_autocast
def computation(res, elem, x, b, t_0_4, t_0_5, s):
  # res: "cpu f32[6]"
  # elem: "cpu f32[6]"
  # xs: "cpu f32[6]"
  # b: "cpu f32[6]"
  # t_0_4: "cpu f32[6]"
  # t_0_5: "cpu f32[6]"
  # s: "cpu f32[6]"
  t0 = torch.add(s, res)  # t0: "cpu f32[6]"
  # ...

These renamings happen in thunder/core/jit_ext._maybe_update_proxy_name, which is called from thunder.core.interpreter._load_fast_handler. frame.code reveals that the irrelevant variable names come from the lookasides implemented in thunder.core.interpreter. res is from SequenceIter.__next__, elem is from _enumerate_lookaside, b from _binary_op.

Ideal behavior

When deciding the identifiers we can just ignore those in thunder's interpreter, which will result in

def computation(x, t_0_1, t_0_2, t_0_3, t_0_4, t_0_5, s):

Altenatively, we can perhaps label the identifier x by an index when x is bound to multiple proxies, as

def computation(x_0, x_1, x_2, x_3, x_4, t_0_5, s):

cc @t-vi @nikitaved

[t-vi]

Very true. Thanks for the writeup and example!

[t-vi]

After #954 we get def computation(x, t_0_1, t_0_2, t_0_3, t_0_4, t_0_5, s): in the example.