[Dynamo][Misc] Apply typing hints for codegen (#150289)

Fixes #ISSUE_NUMBER

Pull Request resolved: https://github.com/pytorch/pytorch/pull/150289
Approved by: https://github.com/Skylion007, https://github.com/cyyever
This commit is contained in:
Yuanhao Ji
2025-04-04 14:26:22 +00:00
committed by PyTorch MergeBot
parent 295b7e21eb
commit 1b0a023dde
17 changed files with 101 additions and 80 deletions

View File

@ -67,6 +67,7 @@ from .user_defined import UserDefinedObjectVariable
if TYPE_CHECKING:
from torch._dynamo.codegen import PyCodegen
from torch._dynamo.symbolic_convert import InstructionTranslator
@ -382,7 +383,7 @@ class TorchFunctionModeVariable(GenericContextWrappingVariable):
self.cm_obj = value # needed for BC with calling enter from CM code
self.source = source
def reconstruct(self, codegen):
def reconstruct(self, codegen: "PyCodegen"):
# This shouldn't be called unless we have a source
assert self.source
self.source.reconstruct(codegen)
@ -426,7 +427,7 @@ class TorchFunctionModeVariable(GenericContextWrappingVariable):
)
return ConstantVariable.create(None)
def reconstruct_type(self, codegen):
def reconstruct_type(self, codegen: "PyCodegen"):
ty = NoEnterTorchFunctionMode
codegen(
AttrSource(