torch.func 的源代码
从 torch._functorch.apis 导入 grad、grad_and_value、vmap,从 torch._functorch.batch_norm_replacement 导入 replace_all_batch_norm_modules_,从 torch._functorch.eager_transforms 导入(debug_unwrap, functionalize, hessian, jacfwd, jacrev, jvp, linearize, vjp, ),从 torch._functorch.functional_call 导入 functional_call, stack_module_state __all__ = ["grad", "grad_and_value", "vmap", "replace_all_batch_norm_modules_", "functionalize", "hessian", "jacfwd", "jacrev", "jvp", "linearize", "vjp", "functional_call", "stack_module_state", "debug_unwrap", ]