Wraps a callable and provide: * tracing via :meth:`.trace` and :meth:`.dump` * accelerated evalutaion via :meth:`.__call__` Args: function(Callable): the function will be traced. symbolic(bool): whether to apply symbolic execution for tracing. Default: False cap
| 93 | |
| 94 | |
| 95 | class trace: |
| 96 | """Wraps a callable and provide: |
| 97 | |
| 98 | * tracing via :meth:`.trace` and :meth:`.dump` |
| 99 | * accelerated evalutaion via :meth:`.__call__` |
| 100 | |
| 101 | Args: |
| 102 | function(Callable): the function will be traced. |
| 103 | symbolic(bool): whether to apply symbolic execution for tracing. Default: False |
| 104 | capture_as_const(bool): capture global vars or closures as const value. Default: False |
| 105 | record_only: if True, won't run even if call the function. Default: False |
| 106 | sublinear_memory_config(SublinearMemoryConfig): configuration for sublinear memory optimization. |
| 107 | If not None, it enables sublinear memory optimization with given setting. |
| 108 | dtr_config(DTRConfig): configuration for DTR sublinear memory optimization. |
| 109 | If not None, it enables DTR optimization with given setting. |
| 110 | profiling(bool): whether to profile compiled trace. Default: False |
| 111 | opt_level(int): optimization level for compiling trace. Default: 2 |
| 112 | graph_opt_config(GraphOptimizationConfig): configuration for graph optimization. Default: None |
| 113 | symbolic_shape(bool): whether to use symbolic shape for tracing. Default: True |
| 114 | without_host(bool): if True, will run python code of wrapped function on the first call, |
| 115 | and run the compiled graph/function on subsequent calls. if False, will run python code every time. |
| 116 | Default: False |
| 117 | imperative(bool): if True, will use imperative runtime to execute captured op seq. Default: False |
| 118 | """ |
| 119 | |
| 120 | third_party_backend = False |
| 121 | |
| 122 | def __new__(cls, *args, **kwargs): |
| 123 | if not args: |
| 124 | return functools.partial(cls, **kwargs) |
| 125 | return super().__new__(cls) |
| 126 | |
| 127 | def __init__( |
| 128 | self, |
| 129 | function, |
| 130 | *, |
| 131 | symbolic=False, |
| 132 | capture_as_const=False, |
| 133 | record_only=False, |
| 134 | sublinear_memory_config: SublinearMemoryConfig = None, |
| 135 | dtr_config: DTRConfig = None, |
| 136 | profiling: bool = False, |
| 137 | opt_level: int = 2, |
| 138 | graph_opt_config: GraphOptimizationConfig = None, |
| 139 | symbolic_shape: bool = True, |
| 140 | without_host: bool = False, |
| 141 | imperative: bool = False, |
| 142 | ): |
| 143 | self.__wrapped__ = function |
| 144 | self._capture_as_const = capture_as_const or record_only |
| 145 | self._arg_bindings = None |
| 146 | self._kwarg_bindings = None |
| 147 | self._output_bindings = None |
| 148 | self._symbolic_shape = symbolic_shape |
| 149 | self._graph_options = { |
| 150 | "no_force_inplace": True, |
| 151 | "graph_opt_level": opt_level, |
| 152 | "seq_opt.enable_seq_comp_node_opt": False, |
no outgoing calls