Compute encoded features. Args: x_input (torch.Tensor): Input tensor (#batch, time, size). mask (torch.Tensor): Mask tensor for the input (#batch, time). cache (torch.Tensor): Cache tensor of the input (#batch, time - 1, size). Returns:
(self, x, cache=None, chunk_size=None, look_back=0)
| 148 | return x, mask, cache, mask_shfit_chunk, mask_att_chunk_encoder |
| 149 | |
| 150 | def forward_chunk(self, x, cache=None, chunk_size=None, look_back=0): |
| 151 | """Compute encoded features. |
| 152 | |
| 153 | Args: |
| 154 | x_input (torch.Tensor): Input tensor (#batch, time, size). |
| 155 | mask (torch.Tensor): Mask tensor for the input (#batch, time). |
| 156 | cache (torch.Tensor): Cache tensor of the input (#batch, time - 1, size). |
| 157 | |
| 158 | Returns: |
| 159 | torch.Tensor: Output tensor (#batch, time, size). |
| 160 | torch.Tensor: Mask tensor (#batch, time). |
| 161 | |
| 162 | """ |
| 163 | |
| 164 | residual = x |
| 165 | if self.normalize_before: |
| 166 | x = self.norm1(x) |
| 167 | |
| 168 | if self.in_size == self.size: |
| 169 | attn, cache = self.self_attn.forward_chunk(x, cache, chunk_size, look_back) |
| 170 | x = residual + attn |
| 171 | else: |
| 172 | x, cache = self.self_attn.forward_chunk(x, cache, chunk_size, look_back) |
| 173 | |
| 174 | if not self.normalize_before: |
| 175 | x = self.norm1(x) |
| 176 | |
| 177 | residual = x |
| 178 | if self.normalize_before: |
| 179 | x = self.norm2(x) |
| 180 | x = residual + self.feed_forward(x) |
| 181 | if not self.normalize_before: |
| 182 | x = self.norm2(x) |
| 183 | |
| 184 | return x, cache |
| 185 | |
| 186 | |
| 187 | @tables.register("encoder_classes", "SANMEncoderChunkOpt") |