| 559 | return hx[0] |
| 560 | |
| 561 | def apply_op(self, input, hx): |
| 562 | fwd_mode = ( |
| 563 | BatchNorm.FwdMode.TRAINING if self.training else BatchNorm.FwdMode.INFERENCE |
| 564 | ) |
| 565 | op = builtin.LSTM( |
| 566 | num_layers=self.num_layers, |
| 567 | bidirectional=self.bidirectional, |
| 568 | bias=self.bias, |
| 569 | hidden_size=self.hidden_size, |
| 570 | proj_size=self.proj_size, |
| 571 | dropout=self.dropout, |
| 572 | fwd_mode=fwd_mode, |
| 573 | ) |
| 574 | output, h, c = apply(op, input, hx[0], hx[1], self._flatten_weights)[:3] |
| 575 | placeholders = [output.sum() * 0, h.sum() * 0, c.sum() * 0] |
| 576 | output = output + placeholders[1] + placeholders[2] |
| 577 | h = h + placeholders[0] + placeholders[2] |
| 578 | c = c + placeholders[0] + placeholders[1] |
| 579 | return output, (h, c) |
| 580 | |
| 581 | def _apply_fn_to_hx(self, hx, fn): |
| 582 | return (fn(hx[0]), fn(hx[1])) |