Shortcuts

functorch.compile.aot_module

functorch.compile.aot_module(mod, *args, **kwargs)[source]

Traces the forward and backward graph of mod using torch dispatch tracing mechanism. It is wrapper function, that underneath uses aot_function() to perform tracing and compilation.

aot_module() lifts the parameters and buffers of nn.Module as inputs to a new callable which is then compiled through aot_function().

Warning

This API is experimental and likely to change.

Parameters
Returns

Returns a nn.Module that retains the eager behavior of the original mod, but with forward and backward graph compiled.

Docs

Access comprehensive developer documentation for PyTorch

View Docs

Tutorials

Get in-depth tutorials for beginners and advanced developers

View Tutorials

Resources

Find development resources and get your questions answered

View Resources