use `torch.OP##_` if available not `torch.Tensor.OP##_`

Lightning-AI / lightning-thunder

Make PyTorch models up to 40% faster! Thunder is a source to source compiler for PyTorch. It enables using different hardware executors at once; across one or thousands of GPUs.

Apache License 2.0

1.15k stars 77 forks source link

use `torch.OP##_` if available not `torch.Tensor.OP##_` #1073

Closed crcrpar closed 3 weeks ago

crcrpar commented 1 month ago

What does this PR do?

Fixes #1072.

import torch

import thunder
from thunder.examine import examine

def f(a):
    return torch.exp_(a)

def g(a):
    return a.exp_()

a = torch.randn(2)
b = thunder.jit(f)(a)
b = thunder.jit(g)(a)

Current main is not capable of functionalizing f, which should be unexpected.