| 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 |
| 152 | |
| 153 | def backward(self, *output_grads): |
| 154 | output_grads = tuple([i for i in output_grads if i is not None]) |
| 155 | return self.bwd(*(output_grads + self.get_keeped_features())) |
| 156 | |
| 157 | class CustomFwd: |
| 158 | def __init__(self, fwd, bwd): |