#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

Extend torch.fx.wrap.

This function can be:

called or used as a decorator on a string to register a builtin : function as a “leaf function”

called or used as a decorator on a function to register this : function as a “leaf function”

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”

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