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