| 76 | |
| 77 | # A causal variant of Conv1d |
| 78 | class CausalConv1d(torch.nn.Conv1d): |
| 79 | def __init__( |
| 80 | self, |
| 81 | in_channels: int, |
| 82 | out_channels: int, |
| 83 | kernel_size: int, |
| 84 | ) -> None: |
| 85 | super(CausalConv1d, self).__init__(in_channels, out_channels, kernel_size) |
| 86 | self.causal_padding = (kernel_size - 1, 0) |
| 87 | |
| 88 | def forward(self, x: torch.Tensor): |
| 89 | x = F.pad(x, self.causal_padding) |
| 90 | x = super(CausalConv1d, self).forward(x) |
| 91 | return x |
| 92 | |
| 93 | def forward_chunk(self, x: torch.Tensor, cnn_cache: torch.Tensor = None): |
| 94 | if cnn_cache is None: |
| 95 | cnn_cache = x.new_zeros( |
| 96 | (x.shape[0], self.in_channels, self.causal_padding[0]) |
| 97 | ) |
| 98 | x = torch.cat([cnn_cache, x], dim=2) |
| 99 | new_cnn_cache = x[..., -self.causal_padding[0] :] |
| 100 | x = super(CausalConv1d, self).forward(x) |
| 101 | return x, new_cnn_cache |
| 102 | |
| 103 | |
| 104 | # A causal variant of ResnetBlock |