x: point features of shape [B, N_feat, 3, N_samples, ...]
(self, x)
| 170 | self.map_to_dir = nn.Linear(in_channels, in_channels, bias=False) |
| 171 | |
| 172 | def forward(self, x): |
| 173 | ''' |
| 174 | x: point features of shape [B, N_feat, 3, N_samples, ...] |
| 175 | ''' |
| 176 | d = self.map_to_dir(x.transpose(1, -1)).transpose(1, -1) |
| 177 | dotprod = (x * d).sum(2, keepdims=True) |
| 178 | idx = dotprod.max(dim=-1, keepdim=False)[1] |
| 179 | index_tuple = torch.meshgrid([torch.arange(j) for j in x.size()[:-1]]) + (idx,) |
| 180 | x_max = x[index_tuple] |
| 181 | return x_max |
| 182 | |
| 183 | |
| 184 | def mean_pool(x, dim=-1, keepdim=False): |
nothing calls this directly
no outgoing calls
no test coverage detected