| 32 | |
| 33 | |
| 34 | class KNNQuery(Function): |
| 35 | @staticmethod |
| 36 | def forward(ctx, nsample, xyz, new_xyz, offset, new_offset): |
| 37 | """ |
| 38 | input: xyz: (n, 3), new_xyz: (m, 3), offset: (b), new_offset: (b) |
| 39 | output: idx: (m, nsample), dist2: (m, nsample) |
| 40 | """ |
| 41 | if new_xyz is None: new_xyz = xyz |
| 42 | assert xyz.is_contiguous() and new_xyz.is_contiguous() |
| 43 | m = new_xyz.shape[0] |
| 44 | idx = torch.cuda.IntTensor(m, nsample).zero_() |
| 45 | dist2 = torch.cuda.FloatTensor(m, nsample).zero_() |
| 46 | pointops_cuda.knnquery_cuda(m, nsample, xyz, new_xyz, offset, new_offset, idx, dist2) |
| 47 | return idx, torch.sqrt(dist2) |
| 48 | |
| 49 | knnquery = KNNQuery.apply |
| 50 |
nothing calls this directly
no outgoing calls
no test coverage detected