(self, X, **kwargs)
| 281 | return Qidxs |
| 282 | |
| 283 | def quantize(self, X, **kwargs): |
| 284 | X = X.T.contiguous().to(torch.float16) |
| 285 | T = X.shape[0] |
| 286 | roll_X = torch.roll(X, T // (2 * self.V) * self.V, 0) |
| 287 | state = self.quantize_seq(roll_X, overlap=None) |
| 288 | overlap = state[T // (2 * self.V)] >> self.K * self.V |
| 289 | state = self.quantize_seq(X, overlap=overlap) |
| 290 | hatX = self.recons(state).transpose(0, 1).reshape(X.shape) |
| 291 | return hatX.T.contiguous().to(X.device), state.T.contiguous().to( |
| 292 | X.device) |
| 293 | |
| 294 | def pack_trellis(self, trellis): |
| 295 | # T is really T // self.V here |
no test coverage detected