Infinite loop in inplace functionalization

[beverlylytle]

The running the following snippet enters an infinite loop:

import thunder
import torch

def foo(a, b):
   x = a.add(1)
   y = x.view(-1)
   z = x.view(-1)
   x.add_(1)
   # y should have been updated.
   y2 = y + 1
   # z should have been updated too.
   z2 = z + 1
   return z2

a = torch.rand(10)
b = torch.rand(10)

jfoo = thunder.(foo)

jfoo(a, b)

Interrupting the program produces the following line in the trace:

  File ".../lightning-thunder/thunder/core/functionalization.py", line 572, in apply_functionalization_to_canonicalized_trace
    cur_v = swap_map[cur_k]

This line comes from a search over the graph defined by swap_map, which terminates when a leaf (a value which is not also a key) is reached. However for the function defined above swap_map contains a key-value pair where the key is the same as the value, resulting in an infinite loop.

Note the behavior is dependent on the shape of a.

[beverlylytle]

The function above is a variation on one appearing in test_multiple_views_before_inplace_to_base.

[crcrpar]

swap_map contains a key-value pair where the key is the same as the value, resulting in an infinite loop.

I find it reasonable to have a check to avoid having a pair of key and value which have the same name.

Note the behavior is dependent on the shape of a.

Would you know a shape that does not cause thunder to get stuck in an infinite loop, out of curiosity?

[beverlylytle]

Would you know a shape that does not cause thunder to get stuck in an infinite loop, out of curiosity?

I didn't investigate this deeply, but 1-dimensional tensors yield an infinite loop, while 2-dimensional tensors do not.