Computing only the n first rows of a distance matrix with R torch

I am writing a fast version of sf::st_is_within_distance(points, dist) for 100k+ points. I start by creating a grid with sf::st_make_grid, then determine which point lies in which grid cell, compute st_is_within_distance(cells, dist) and then loop on the grid cells. For each grid cell, I need to compute the distances between the points in the current cell and the points in cells within dist of the current cell (including the current cell itself).

With p_current the set of points in the current cell, and p_neighbors the set of points in the cells within dist of the current cell (notice that, if ordered correctly, p_neighbors starts with p_current), I have tried:

• run torch::pdist(p_neighbors), and throw away most lines of the upper-triangular matrix it returns. It works but is a huge waste
• run torch::pairwise_distance(v1, v2), where v1 = p_neighbors[row_indices] and v2 = p_neighbors[columns_indices], with row_indices and column_indices the indices of the elements of the first rows of an upper triangular matrix, diagonal excluded. It also works, but creating v1 and v2 results in a significant overhead

Is there a way to call torch::pdist(p_neighbors), but ask it to compute and return only the n first rows of the distance matrix ? If not, would it be difficult to add, would it be relevant, and more importantly, where should I ask about it ?

Currently, PyTorch's pdist() doesn't have a built-in function to compute only specific rows of the distance matrix. This is PyTorch GitHub repository: https://github.com/pytorch/pytorch/issues this would be a better place to discuss and find a solution