(self, x: SparseTensor, subdivision: Optional[SparseTensor] = None)
| 80 | self.factor = factor |
| 81 | |
| 82 | def forward(self, x: SparseTensor, subdivision: Optional[SparseTensor] = None) -> SparseTensor: |
| 83 | DIM = x.coords.shape[-1] - 1 |
| 84 | |
| 85 | cache = x.get_spatial_cache(f'upsample_{self.factor}') |
| 86 | if cache is None: |
| 87 | if subdivision is None: |
| 88 | raise ValueError('Cache not found. Provide subdivision tensor or pair SparseUpsample with SparseDownsample.') |
| 89 | else: |
| 90 | sub = subdivision.feats |
| 91 | N_leaf = sub.sum(dim=-1) |
| 92 | subidx = sub.nonzero()[:, -1] |
| 93 | new_coords = x.coords.clone().detach() |
| 94 | new_coords[:, 1:] *= self.factor |
| 95 | new_coords = torch.repeat_interleave(new_coords, N_leaf, dim=0, output_size=subidx.shape[0]) |
| 96 | for i in range(DIM): |
| 97 | new_coords[:, i+1] += subidx // self.factor ** i % self.factor |
| 98 | idx = torch.repeat_interleave(torch.arange(x.coords.shape[0], device=x.device), N_leaf, dim=0, output_size=subidx.shape[0]) |
| 99 | else: |
| 100 | new_coords, idx = cache |
| 101 | |
| 102 | new_feats = x.feats[idx] |
| 103 | out = SparseTensor(new_feats, new_coords, x._shape) |
| 104 | out._scale = tuple([s / self.factor for s in x._scale]) |
| 105 | if cache is not None: # only keep cache when subdiv following it |
| 106 | out._spatial_cache = x._spatial_cache |
| 107 | |
| 108 | return out |
| 109 |
nothing calls this directly
no test coverage detected