| from torch._functorch.python_key import pythonkey_decompose | |
| from torch._functorch.fx_minifier import minifier | |
| from torch._functorch.aot_autograd import ( | |
| aot_function, | |
| aot_module, | |
| compiled_function, | |
| compiled_module, | |
| aot_module_simplified, | |
| get_graph_being_compiled, | |
| get_aot_graph_name, | |
| get_aot_compilation_context, | |
| make_boxed_func, | |
| make_boxed_compiler | |
| ) | |
| from torch._functorch.compilers import ( | |
| ts_compile, | |
| draw_graph_compile, | |
| nop, | |
| nnc_jit, | |
| memory_efficient_fusion, | |
| debug_compile, | |
| print_compile, | |
| default_decompositions | |
| ) | |
| from torch._functorch.partitioners import ( | |
| min_cut_rematerialization_partition, | |
| default_partition, | |
| draw_graph, | |
| draw_joint_graph, | |
| ) | |
| from torch._functorch import config | |