__init__.py 776 B

12345678910111213141516171819202122232425262728293031
  1. from torch._functorch.python_key import pythonkey_decompose
  2. from torch._functorch.fx_minifier import minifier
  3. from torch._functorch.aot_autograd import (
  4. aot_function,
  5. aot_module,
  6. compiled_function,
  7. compiled_module,
  8. aot_module_simplified,
  9. get_graph_being_compiled,
  10. get_aot_graph_name,
  11. get_aot_compilation_context,
  12. make_boxed_func,
  13. make_boxed_compiler
  14. )
  15. from torch._functorch.compilers import (
  16. ts_compile,
  17. draw_graph_compile,
  18. nop,
  19. nnc_jit,
  20. memory_efficient_fusion,
  21. debug_compile,
  22. print_compile,
  23. default_decompositions
  24. )
  25. from torch._functorch.partitioners import (
  26. min_cut_rematerialization_partition,
  27. default_partition,
  28. draw_graph,
  29. draw_joint_graph,
  30. )
  31. from torch._functorch import config