[visualization] update make_trace_dot

[kshitij12345]

Builds on great feature introduced by @riccardofelluga

Changes

• We add a seperate node for input/output to a Symbol (with option to add more metadata)
• Symbol node prints the python repr to show more information (currently we print the symbol name)
• Update color scheme - Roots nodes are colored "green", intermediates are colored "lightblue" and leaves are colored "orange"

import thunder
from thunder.examine import make_trace_dot
import torch

def foo(x, y):
    return x + y / (x * y)

jfoo = thunder.jit(foo)

jfoo(torch.randn(3, requires_grad=True), torch.randn(3,))

t = thunder.last_traces(jfoo)[-3]
dot = make_trace_dot(t)
dot.render(format='png')

Output

Forward Graph

Backward Graph

[riccardofelluga]

Minor things but overall it's a great addition!

[riccardofelluga]

It's a good idea to leave the rendering to the user at they will have the dot object anyway. Stamping! 🎉

[t-vi]

Great, thank you @kshitij12345 @riccardofelluga