#fx.fx_helper.wrap

Extended tracer and wrap of torch.fx.

This file defines a inherit tracer of torch.fx.Tracer and a extended wrap to allow wrapping of user-defined Module or method, which help users do some optimization of their own module by torch.fx

#horizon_plugin_pytorch.fx.fx_helper.wrap(skip_compile: bool = False)

Extend torch.fx.wrap.

This function can be: : 1) called or used as a decorator on a string to register a builtin function as a “leaf function” 2) called or used as a decorator on a function to register this function as a “leaf function” 3) called or used as a decorator on subclass of torch.nn.Module to register this module as a “leaf module”, and register all user defined method in this class as “leaf method” 4) called or used as a decorator on a class method to register it as “leaf method”

• Parameters: skip_compile (bool , optional) – Whether wrapped obj is skipped in compile, used by horizon_plugin_pytorch.quantization.fx.split_compilable_model .split_compilable_model. Defaults to False.
• Returns: The actural decorator.
• Return type: FxWrapManager.wrap