[autograd.Function] freevar lifting is too aggressive?

[zou3519]

🐛 Describe the bug

Seems like the freevar lifting things it needs to lift constants and then graph breaks if they are used in autograd.Function backwards:

import torch
from typing import *

def halve(x):
    return x * 0.5


class ScaleGradient(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        return x

    @staticmethod
    def backward(ctx, grad):
        return halve(x)

x = torch.randn(3, requires_grad=True)
def f(x):
    return ScaleGradient.apply(x)
output = torch.compile(f, backend='eager', fullgraph=True)(x)

gives:

    f.call_function(tx, sub_args, {})
  File "/raid/rzou/pt/debug-cpu2/torch/_dynamo/variables/higher_order_ops.py", line 950, in call_funct
ion
    unimplemented("NYI - freevars in autograd function.")
  File "/raid/rzou/pt/debug-cpu2/torch/_dynamo/exc.py", line 143, in unimplemented
    raise Unsupported(msg)
torch._dynamo.exc.Unsupported: NYI - freevars in autograd function.

from user code:
   File "/raid/rzou/pt/debug-cpu2/foo.py", line 20, in f
    return ScaleGradient.apply(x)

Set TORCH_LOGS="+dynamo" and TORCHDYNAMO_VERBOSE=1 for more information


You can suppress this exception and fall back to eager by setting:
    import torch._dynamo
    torch._dynamo.config.suppress_errors = True

cc @ezyang @anijain2305 @chauhang @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng @ydwu4 @msaroufim @bdhirsh @aakhundov @wconstab @Xia-Weiwen

Error logs

No response

Minified repro

No response

Versions

main

[zou3519]

Succeeds now