(self, splats: TensorDict, quant_ctx: QuantizationContext | None = None)
| 61 | return trimmed |
| 62 | |
| 63 | def map(self, splats: TensorDict, quant_ctx: QuantizationContext | None = None) -> Tuple[TensorDict, MappingContext]: |
| 64 | if not self.config.enabled or self.config.strategy == "none": |
| 65 | return {k: v.clone() for k, v in splats.items()}, MappingContext(strategy="none", permutation=None) |
| 66 | |
| 67 | working = {k: v.clone() for k, v in splats.items()} |
| 68 | if quant_ctx is None: |
| 69 | quant_ctx = QuantizationContext() |
| 70 | |
| 71 | working = self._ensure_square_length(working, quant_ctx) |
| 72 | |
| 73 | if self.config.strategy == "morton": |
| 74 | mapped, perm = sort_splats_morton(working, verbose=self.config.verbose, return_indices=True) |
| 75 | # apply permutation to cached integer tensors used by codecs |
| 76 | for name, int_tensor in list(quant_ctx.int_values.items()): |
| 77 | if int_tensor.shape[0] == perm.shape[0]: |
| 78 | quant_ctx.int_values[name] = int_tensor[perm] |
| 79 | # also permute any per-point masks stored in stats (for VQ) |
| 80 | perm_cpu = perm.cpu() |
| 81 | for fname, stats in quant_ctx.field_stats.items(): |
| 82 | if stats.mask is not None and len(stats.mask) == perm_cpu.shape[0]: |
| 83 | mask_tensor = torch.tensor(stats.mask, dtype=torch.int8) |
| 84 | stats.mask = mask_tensor[perm_cpu].to(torch.int32).tolist() |
| 85 | return mapped, MappingContext(strategy="morton", permutation=perm) |
| 86 | if self.config.strategy == "plas": |
| 87 | mapped, perm = sort_splats( |
| 88 | working, |
| 89 | verbose=self.config.verbose, |
| 90 | return_indices=True, |
| 91 | sort_with_shN=self.config.sort_with_shN, |
| 92 | ) |
| 93 | # apply permutation to cached integer tensors used by codecs |
| 94 | for name, int_tensor in list(quant_ctx.int_values.items()): |
| 95 | if int_tensor.shape[0] == perm.shape[0]: |
| 96 | quant_ctx.int_values[name] = int_tensor[perm] |
| 97 | # also permute any per-point masks stored in stats (for VQ) |
| 98 | perm_cpu = perm.cpu() |
| 99 | for fname, stats in quant_ctx.field_stats.items(): |
| 100 | if stats.mask is not None and len(stats.mask) == perm_cpu.shape[0]: |
| 101 | mask_tensor = torch.tensor(stats.mask, dtype=torch.int8) |
| 102 | stats.mask = mask_tensor[perm_cpu].to(torch.int32).tolist() |
| 103 | return mapped, MappingContext(strategy="plas", permutation=perm) |
| 104 | return working, MappingContext(strategy=self.config.strategy, permutation=None) |
no test coverage detected