| 360 | return quantized_out, out_indices, out_losses, out_quantized |
| 361 | |
| 362 | def encode( |
| 363 | self, x: torch.Tensor, n_q: tp.Optional[int] = None, st: tp.Optional[int] = None |
| 364 | ) -> torch.Tensor: |
| 365 | residual = x |
| 366 | all_indices = [] |
| 367 | n_q = n_q or len(self.layers) |
| 368 | st = st or 0 |
| 369 | for layer in self.layers[st:n_q]: |
| 370 | indices = layer.encode(residual) |
| 371 | quantized = layer.decode(indices) |
| 372 | residual = residual - quantized |
| 373 | all_indices.append(indices) |
| 374 | out_indices = torch.stack(all_indices) |
| 375 | return out_indices |
| 376 | |
| 377 | def decode(self, q_indices: torch.Tensor, st: int = 0) -> torch.Tensor: |
| 378 | quantized_out = torch.tensor(0.0, device=q_indices.device) |