mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-27 17:54:55 +08:00
cce2c52b0b
Teach dynamo about `vmap` Pull Request resolved: https://github.com/pytorch/pytorch/pull/101707 Approved by: https://github.com/zou3519
15 lines
401 B
Python
15 lines
401 B
Python
from torch._functorch.eager_transforms import (
|
|
grad_and_value,
|
|
vjp,
|
|
jvp,
|
|
jacrev,
|
|
jacfwd,
|
|
hessian,
|
|
functionalize,
|
|
linearize
|
|
)
|
|
from torch._functorch.apis import grad
|
|
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
|