| 92 | |
| 93 | |
| 94 | def partial_trace( |
| 95 | func=None, |
| 96 | *, |
| 97 | backend="default", |
| 98 | without_host=True, |
| 99 | trace_gmcallback=True, |
| 100 | **trace_options |
| 101 | ): |
| 102 | |
| 103 | assert backend in JIT_BACKEND |
| 104 | assert without_host, "partial_trace only support without_host mode currently!" |
| 105 | |
| 106 | def wrapper(func): |
| 107 | trace_obj = JIT_BACKEND[backend]( |
| 108 | func, without_host=without_host, **trace_options |
| 109 | ) |
| 110 | trace_options["capture_as_const"] = False |
| 111 | backward_trace_obj = JIT_BACKEND[backend]( |
| 112 | None, without_host=without_host, **trace_options |
| 113 | ) |
| 114 | backward_trace_obj.check_external = ( |
| 115 | False # check if there are unknown external vars after tracing. |
| 116 | ) |
| 117 | trace_obj.overall = False # if trace overall train step |
| 118 | backward_trace_obj.overall = False |
| 119 | trace_obj._trace.remove_unused_data_required = False |
| 120 | backward_trace_obj._trace.remove_unused_data_required = False |
| 121 | inp_grad_maps = OrderedDict() # x, dx map |
| 122 | out_grad_maps = OrderedDict() # y, dy map |
| 123 | traced = False # if wrapped function has been traced |
| 124 | custom_autodiff = None |
| 125 | outdef = None # treedef of forward return value |
| 126 | check_shape = backend == "xla" |
| 127 | shape_hash = None |
| 128 | unexpected_shape_hash = set() |
| 129 | from ..core.autodiff.grad import Function |
| 130 | |
| 131 | class CustomAutodiff(Function): |
| 132 | def __init__(self, fwd, bwd): |
| 133 | self.fwd = fwd |
| 134 | self.bwd = bwd |
| 135 | del fwd.outdef |
| 136 | self.keeped_features = [] |
| 137 | |
| 138 | def forward(self, *args): |
| 139 | rst = self.fwd(*args) |
| 140 | keeped_features = rst[-1] |
| 141 | if not isinstance(keeped_features, Sequence): |
| 142 | keeped_features = tuple([keeped_features]) |
| 143 | else: |
| 144 | keeped_features = tuple(keeped_features) |
| 145 | self.keeped_features = keeped_features |
| 146 | return rst[0] |
| 147 | |
| 148 | def get_keeped_features(self): |
| 149 | rst = self.keeped_features |
| 150 | del self.keeped_features |
| 151 | return rst |