(x, y, k, batch_x=None, batch_y=None)
| 41 | |
| 42 | |
| 43 | def knn_atoms(x, y, k, batch_x=None, batch_y=None): |
| 44 | k = k+1 |
| 45 | N, D = x.shape |
| 46 | x_i = LazyTensor(x[:, None, :]) |
| 47 | y_j = LazyTensor(y[None, :, :]) |
| 48 | |
| 49 | pairwise_distance_ij = ((x_i - y_j) ** 2).sum(-1) |
| 50 | if batch_x is not None: |
| 51 | pairwise_distance_ij.ranges = diagonal_ranges(batch_x, batch_y) |
| 52 | |
| 53 | # N.B.: KeOps doesn't yet support backprop through Kmin reductions... |
| 54 | # dists, idx = pairwise_distance_ij.Kmin_argKmin(K=k,axis=1) |
| 55 | # So we have to re-compute the values ourselves: |
| 56 | idx = pairwise_distance_ij.argKmin(K=k, axis=1) # (N, K) |
| 57 | x_ik = y[idx.view(-1)].view(N, k, D) |
| 58 | dists = ((x[:, None, :] - x_ik) ** 2).sum(-1) |
| 59 | |
| 60 | return idx, dists |
| 61 | |
| 62 | |
| 63 | def compute_eigens(n_verts, xyz, min_n_eigs, eigs_ratio): |
nothing calls this directly
no test coverage detected