| 155 | return self.bwd(*(output_grads + self.get_keeped_features())) |
| 156 | |
| 157 | class CustomFwd: |
| 158 | def __init__(self, fwd, bwd): |
| 159 | self.fwd = fwd |
| 160 | self.bwd = bwd |
| 161 | |
| 162 | def __call__(self, *args): |
| 163 | rst = self.fwd(*args) |
| 164 | if self.fwd.keeped_activation: |
| 165 | keeped_features = rst[-1] |
| 166 | if not isinstance(keeped_features, Sequence): |
| 167 | keeped_features = tuple([keeped_features]) |
| 168 | else: |
| 169 | keeped_features = tuple(keeped_features) |
| 170 | self.keeped_features = keeped_features |
| 171 | return rst[0] |
| 172 | else: |
| 173 | return rst |
| 174 | |
| 175 | def get_shape_hash(*tensors): |
| 176 | def map_scalar_to_tuple(ishape): |