Files
pytorch/torch/func/__init__.py
Guilherme Leobas 491c2b4665 Let torch dynamo inline torch.func.grad (#118407)
When dynamo sees torch.func.grad, it tries to inline all frames related
to.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/118407
Approved by: https://github.com/zou3519
2024-02-28 20:05:00 +00:00

14 lines
397 B
Python

from torch._functorch.eager_transforms import (
vjp,
jvp,
jacrev,
jacfwd,
hessian,
functionalize,
linearize
)
from torch._functorch.apis import grad, grad_and_value
from torch._functorch.functional_call import functional_call, stack_module_state
from torch._functorch.batch_norm_replacement import replace_all_batch_norm_modules_
from torch._functorch.apis import vmap