(self, input: SparseTensor)
| 90 | super(SparseSubdivide, self).__init__() |
| 91 | |
| 92 | def forward(self, input: SparseTensor) -> SparseTensor: |
| 93 | DIM = input.coords.shape[-1] - 1 |
| 94 | # upsample scale=2^DIM |
| 95 | n_cube = torch.ones([2] * DIM, device=input.device, dtype=torch.int) |
| 96 | n_coords = torch.nonzero(n_cube) |
| 97 | n_coords = torch.cat([torch.zeros_like(n_coords[:, :1]), n_coords], dim=-1) |
| 98 | factor = n_coords.shape[0] |
| 99 | assert factor == 2 ** DIM |
| 100 | # print(n_coords.shape) |
| 101 | new_coords = input.coords.clone() |
| 102 | new_coords[:, 1:] *= 2 |
| 103 | new_coords = new_coords.unsqueeze(1) + n_coords.unsqueeze(0).to(new_coords.dtype) |
| 104 | |
| 105 | new_feats = input.feats.unsqueeze(1).expand(input.feats.shape[0], factor, *input.feats.shape[1:]) |
| 106 | out = SparseTensor(new_feats.flatten(0, 1), new_coords.flatten(0, 1), input.shape) |
| 107 | out._scale = input._scale * 2 |
| 108 | out._spatial_cache = input._spatial_cache |
| 109 | return out |
| 110 |
nothing calls this directly
no test coverage detected