fix memory leak in knn/radius

csrc/cpu/knn_cpu.cpp

 #include "knn_cpu.h"
 
 #include "utils.h"
-#include "utils/neighbors.cpp"
+#include "utils/KDTreeVectorOfVectorsAdaptor.h"
+#include "utils/nanoflann.hpp"
 
 torch::Tensor knn_cpu(torch::Tensor x, torch::Tensor y,
                       torch::optional<torch::Tensor> ptr_x,
@@ -25,25 +26,72 @@ torch::Tensor knn_cpu(torch::Tensor x, torch::Tensor y,
   std::vector<size_t> out_vec = std::vector<size_t>();
 
   AT_DISPATCH_ALL_TYPES(x.scalar_type(), "knn_cpu", [&] {
+    // See: nanoflann/examples/vector_of_vectors_example.cpp
+
     auto x_data = x.data_ptr<scalar_t>();
     auto y_data = y.data_ptr<scalar_t>();
-    auto x_vec = std::vector<scalar_t>(x_data, x_data + x.numel());
-    auto y_vec = std::vector<scalar_t>(y_data, y_data + y.numel());
-
-    if (!ptr_x.has_value()) {
-      nanoflann_neighbors<scalar_t>(y_vec, x_vec, out_vec, 0, x.size(-1), 0,
-                                    num_workers, k, 0);
-    } else {
-      auto sx = (ptr_x.value().narrow(0, 1, ptr_x.value().numel() - 1) -
-                 ptr_x.value().narrow(0, 0, ptr_x.value().numel() - 1));
-      auto sy = (ptr_y.value().narrow(0, 1, ptr_y.value().numel() - 1) -
-                 ptr_y.value().narrow(0, 0, ptr_y.value().numel() - 1));
-      auto sx_data = sx.data_ptr<int64_t>();
-      auto sy_data = sy.data_ptr<int64_t>();
-      auto sx_vec = std::vector<long>(sx_data, sx_data + sx.numel());
-      auto sy_vec = std::vector<long>(sy_data, sy_data + sy.numel());
-      batch_nanoflann_neighbors<scalar_t>(y_vec, x_vec, sy_vec, sx_vec, out_vec,
-                                          k, x.size(-1), 0, k, 0);
+    typedef std::vector<std::vector<scalar_t>> vec_t;
+
+    if (!ptr_x.has_value()) { // Single example.
+
+      vec_t pts(x.size(0));
+      for (int64_t i = 0; i < x.size(0); i++) {
+        pts[i].resize(x.size(1));
+        for (int64_t j = 0; j < x.size(1); j++) {
+          pts[i][j] = x_data[i * x.size(1) + j];
+        }
+      }
+
+      typedef KDTreeVectorOfVectorsAdaptor<vec_t, scalar_t> my_kd_tree_t;
+
+      my_kd_tree_t mat_index(x.size(1), pts, 10);
+      mat_index.index->buildIndex();
+
+      std::vector<size_t> ret_index(k);
+      std::vector<scalar_t> out_dist_sqr(k);
+      for (int64_t i = 0; i < y.size(0); i++) {
+        size_t num_matches = mat_index.index->knnSearch(
+            y_data + i * y.size(1), k, &ret_index[0], &out_dist_sqr[0]);
+
+        for (size_t j = 0; j < num_matches; j++) {
+          out_vec.push_back(ret_index[j]);
+          out_vec.push_back(i);
+        }
+      }
+    } else { // Batch-wise.
+
+      auto ptr_x_data = ptr_x.value().data_ptr<int64_t>();
+      auto ptr_y_data = ptr_y.value().data_ptr<int64_t>();
+
+      for (int64_t b = 0; b < ptr_x.value().size(0) - 1; b++) {
+        auto x_start = ptr_x_data[b], x_end = ptr_x_data[b + 1];
+        auto y_start = ptr_y_data[b], y_end = ptr_y_data[b + 1];
+
+        vec_t pts(x_end - x_start);
+        for (int64_t i = 0; i < x_end - x_start; i++) {
+          pts[i].resize(x.size(1));
+          for (int64_t j = 0; j < x.size(1); j++) {
+            pts[i][j] = x_data[(i + x_start) * x.size(1) + j];
+          }
+        }
+
+        typedef KDTreeVectorOfVectorsAdaptor<vec_t, scalar_t> my_kd_tree_t;
+
+        my_kd_tree_t mat_index(x.size(1), pts, 10);
+        mat_index.index->buildIndex();
+
+        std::vector<size_t> ret_index(k);
+        std::vector<scalar_t> out_dist_sqr(k);
+        for (int64_t i = y_start; i < y_end; i++) {
+          size_t num_matches = mat_index.index->knnSearch(
+              y_data + i * y.size(1), k, &ret_index[0], &out_dist_sqr[0]);
+
+          for (size_t j = 0; j < num_matches; j++) {
+            out_vec.push_back(x_start + ret_index[j]);
+            out_vec.push_back(i);
+          }
+        }
+      }
     }
   });
 

csrc/cpu/radius_cpu.cpp

 #include "radius_cpu.h"
 
 #include "utils.h"
-#include "utils/neighbors.cpp"
+#include "utils/KDTreeVectorOfVectorsAdaptor.h"
+#include "utils/nanoflann.hpp"
 
 torch::Tensor radius_cpu(torch::Tensor x, torch::Tensor y,
                          torch::optional<torch::Tensor> ptr_x,
@@ -25,26 +26,74 @@ torch::Tensor radius_cpu(torch::Tensor x, torch::Tensor y,
   std::vector<size_t> out_vec = std::vector<size_t>();
 
   AT_DISPATCH_ALL_TYPES(x.scalar_type(), "radius_cpu", [&] {
+    // See: nanoflann/examples/vector_of_vectors_example.cpp
+
     auto x_data = x.data_ptr<scalar_t>();
     auto y_data = y.data_ptr<scalar_t>();
-    auto x_vec = std::vector<scalar_t>(x_data, x_data + x.numel());
-    auto y_vec = std::vector<scalar_t>(y_data, y_data + y.numel());
-
-    if (!ptr_x.has_value()) {
-      nanoflann_neighbors<scalar_t>(y_vec, x_vec, out_vec, r, x.size(-1),
-                                    max_num_neighbors, num_workers, 0, 1);
-    } else {
-      auto sx = (ptr_x.value().narrow(0, 1, ptr_x.value().numel() - 1) -
-                 ptr_x.value().narrow(0, 0, ptr_x.value().numel() - 1));
-      auto sy = (ptr_y.value().narrow(0, 1, ptr_y.value().numel() - 1) -
-                 ptr_y.value().narrow(0, 0, ptr_y.value().numel() - 1));
-      auto sx_data = sx.data_ptr<int64_t>();
-      auto sy_data = sy.data_ptr<int64_t>();
-      auto sx_vec = std::vector<long>(sx_data, sx_data + sx.numel());
-      auto sy_vec = std::vector<long>(sy_data, sy_data + sy.numel());
-      batch_nanoflann_neighbors<scalar_t>(y_vec, x_vec, sy_vec, sx_vec, out_vec,
-                                          r, x.size(-1), max_num_neighbors, 0,
-                                          1);
+    typedef std::vector<std::vector<scalar_t>> vec_t;
+    nanoflann::SearchParams params;
+    params.sorted = false;
+
+    if (!ptr_x.has_value()) { // Single example.
+
+      vec_t pts(x.size(0));
+      for (int64_t i = 0; i < x.size(0); i++) {
+        pts[i].resize(x.size(1));
+        for (int64_t j = 0; j < x.size(1); j++) {
+          pts[i][j] = x_data[i * x.size(1) + j];
+        }
+      }
+
+      typedef KDTreeVectorOfVectorsAdaptor<vec_t, scalar_t> my_kd_tree_t;
+
+      my_kd_tree_t mat_index(x.size(1), pts, 10);
+      mat_index.index->buildIndex();
+
+      for (int64_t i = 0; i < y.size(0); i++) {
+        std::vector<std::pair<size_t, scalar_t>> ret_matches;
+        size_t num_matches = mat_index.index->radiusSearch(
+            y_data + i * y.size(1), r * r + 0.00001, ret_matches, params);
+
+        for (size_t j = 0; j < std::min(num_matches, (size_t)max_num_neighbors);
+             j++) {
+          out_vec.push_back(ret_matches[j].first);
+          out_vec.push_back(i);
+        }
+      }
+
+    } else { // Batch-wise.
+
+      auto ptr_x_data = ptr_x.value().data_ptr<int64_t>();
+      auto ptr_y_data = ptr_y.value().data_ptr<int64_t>();
+
+      for (int64_t b = 0; b < ptr_x.value().size(0) - 1; b++) {
+        auto x_start = ptr_x_data[b], x_end = ptr_x_data[b + 1];
+        auto y_start = ptr_y_data[b], y_end = ptr_y_data[b + 1];
+
+        vec_t pts(x_end - x_start);
+        for (int64_t i = 0; i < x_end - x_start; i++) {
+          pts[i].resize(x.size(1));
+          for (int64_t j = 0; j < x.size(1); j++) {
+            pts[i][j] = x_data[(i + x_start) * x.size(1) + j];
+          }
+        }
+
+        typedef KDTreeVectorOfVectorsAdaptor<vec_t, scalar_t> my_kd_tree_t;
+
+        my_kd_tree_t mat_index(x.size(1), pts, 10);
+        mat_index.index->buildIndex();
+
+        for (int64_t i = y_start; i < y_end; i++) {
+          std::vector<std::pair<size_t, scalar_t>> ret_matches;
+          size_t num_matches = mat_index.index->radiusSearch(
+              y_data + i * y.size(1), r * r + 0.00001, ret_matches, params);
+
+          for (size_t j = 0; j < num_matches; j++) {
+            out_vec.push_back(x_start + ret_matches[j].first);
+            out_vec.push_back(i);
+          }
+        }
+      }
     }
   });
 

csrc/cpu/utils/KDTreeVectorOfVectorsAdaptor.h

+/***********************************************************************
+ * Software License Agreement (BSD License)
+ *
+ * Copyright 2011-16 Jose Luis Blanco (joseluisblancoc@gmail.com).
+ *   All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions
+ * are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright
+ *    notice, this list of conditions and the following disclaimer.
+ * 2. Redistributions in binary form must reproduce the above copyright
+ *    notice, this list of conditions and the following disclaimer in the
+ *    documentation and/or other materials provided with the distribution.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR
+ * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
+ * OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
+ * IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT,
+ * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT
+ * NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+ * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+ * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
+ * THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *************************************************************************/
+
+#pragma once
+
+#include "nanoflann.hpp"
+
+#include <vector>
+
+// ===== This example shows how to use nanoflann with these types of containers:
+// =======
+// typedef std::vector<std::vector<double> > my_vector_of_vectors_t;
+// typedef std::vector<Eigen::VectorXd> my_vector_of_vectors_t;   // This
+// requires #include <Eigen/Dense>
+// =====================================================================================
+
+/** A simple vector-of-vectors adaptor for nanoflann, without duplicating the
+ * storage. The i'th vector represents a point in the state space.
+ *
+ *  \tparam DIM If set to >0, it specifies a compile-time fixed dimensionality
+ * for the points in the data set, allowing more compiler optimizations. \tparam
+ * num_t The type of the point coordinates (typically, double or float). \tparam
+ * Distance The distance metric to use: nanoflann::metric_L1,
+ * nanoflann::metric_L2, nanoflann::metric_L2_Simple, etc. \tparam IndexType The
+ * type for indices in the KD-tree index (typically, size_t of int)
+ */
+template <class VectorOfVectorsType, typename num_t = double, int DIM = -1,
+          class Distance = nanoflann::metric_L2, typename IndexType = size_t>
+struct KDTreeVectorOfVectorsAdaptor {
+  typedef KDTreeVectorOfVectorsAdaptor<VectorOfVectorsType, num_t, DIM,
+                                       Distance>
+      self_t;
+  typedef
+      typename Distance::template traits<num_t, self_t>::distance_t metric_t;
+  typedef nanoflann::KDTreeSingleIndexAdaptor<metric_t, self_t, DIM, IndexType>
+      index_t;
+
+  index_t *index; //! The kd-tree index for the user to call its methods as
+                  //! usual with any other FLANN index.
+
+  /// Constructor: takes a const ref to the vector of vectors object with the
+  /// data points
+  KDTreeVectorOfVectorsAdaptor(const size_t /* dimensionality */,
+                               const VectorOfVectorsType &mat,
+                               const int leaf_max_size = 10)
+      : m_data(mat) {
+    assert(mat.size() != 0 && mat[0].size() != 0);
+    const size_t dims = mat[0].size();
+    if (DIM > 0 && static_cast<int>(dims) != DIM)
+      throw std::runtime_error(
+          "Data set dimensionality does not match the 'DIM' template argument");
+    index =
+        new index_t(static_cast<int>(dims), *this /* adaptor */,
+                    nanoflann::KDTreeSingleIndexAdaptorParams(leaf_max_size));
+    index->buildIndex();
+  }
+
+  ~KDTreeVectorOfVectorsAdaptor() { delete index; }
+
+  const VectorOfVectorsType &m_data;
+
+  /** Query for the \a num_closest closest points to a given point (entered as
+   * query_point[0:dim-1]). Note that this is a short-cut method for
+   * index->findNeighbors(). The user can also call index->... methods as
+   * desired. \note nChecks_IGNORED is ignored but kept for compatibility with
+   * the original FLANN interface.
+   */
+  inline void query(const num_t *query_point, const size_t num_closest,
+                    IndexType *out_indices, num_t *out_distances_sq,
+                    const int nChecks_IGNORED = 10) const {
+    nanoflann::KNNResultSet<num_t, IndexType> resultSet(num_closest);
+    resultSet.init(out_indices, out_distances_sq);
+    index->findNeighbors(resultSet, query_point, nanoflann::SearchParams());
+  }
+
+  /** @name Interface expected by KDTreeSingleIndexAdaptor
+   * @{ */
+
+  const self_t &derived() const { return *this; }
+  self_t &derived() { return *this; }
+
+  // Must return the number of data points
+  inline size_t kdtree_get_point_count() const { return m_data.size(); }
+
+  // Returns the dim'th component of the idx'th point in the class:
+  inline num_t kdtree_get_pt(const size_t idx, const size_t dim) const {
+    return m_data[idx][dim];
+  }
+
+  // Optional bounding-box computation: return false to default to a standard
+  // bbox computation loop.
+  //   Return true if the BBOX was already computed by the class and returned in
+  //   "bb" so it can be avoided to redo it again. Look at bb.size() to find out
+  //   the expected dimensionality (e.g. 2 or 3 for point clouds)
+  template <class BBOX> bool kdtree_get_bbox(BBOX & /*bb*/) const {
+    return false;
+  }
+
+  /** @} */
+
+}; // end of KDTreeVectorOfVectorsAdaptor

csrc/cpu/utils/cloud.h

-#pragma once
-
-#include <ATen/ATen.h>
-#include <algorithm>
-#include <cmath>
-#include <iomanip>
-#include <iostream>
-#include <map>
-#include <numeric>
-#include <unordered_map>
-#include <vector>
-
-#include <time.h>
-
-template <typename scalar_t> struct PointCloud {
-  std::vector<std::vector<scalar_t> *> pts;
-
-  void set(std::vector<scalar_t> new_pts, int dim) {
-
-    std::vector<std::vector<scalar_t> *> temp(new_pts.size() / dim);
-    for (size_t i = 0; i < new_pts.size(); i++) {
-      if (i % dim == 0) {
-        std::vector<scalar_t> *point = new std::vector<scalar_t>(dim);
-
-        for (size_t j = 0; j < (size_t)dim; j++) {
-          (*point)[j] = new_pts[i + j];
-        }
-        temp[i / dim] = point;
-      }
-    }
-
-    pts = temp;
-  }
-  void set_batch(std::vector<scalar_t> new_pts, size_t begin, long size,
-                 int dim) {
-    std::vector<std::vector<scalar_t> *> temp(size);
-    for (size_t i = 0; i < (size_t)size; i++) {
-      std::vector<scalar_t> *point = new std::vector<scalar_t>(dim);
-      for (size_t j = 0; j < (size_t)dim; j++) {
-        (*point)[j] = new_pts[dim * (begin + i) + j];
-      }
-
-      temp[i] = point;
-    }
-    pts = temp;
-  }
-
-  // Must return the number of data points.
-  inline size_t kdtree_get_point_count() const { return pts.size(); }
-
-  // Returns the dim'th component of the idx'th point in the class:
-  inline scalar_t kdtree_get_pt(const size_t idx, const size_t dim) const {
-    return (*pts[idx])[dim];
-  }
-
-  // Optional bounding-box computation: return false to default to a standard
-  // bbox computation loop.
-  // Return true if the BBOX was already computed by the class and returned in
-  // "bb" so it can be avoided to redo it again. Look at bb.size() to find out
-  // the expected dimensionality (e.g. 2 or 3 for point clouds)
-  template <class BBOX> bool kdtree_get_bbox(BBOX & /* bb */) const {
-    return false;
-  }
-};

csrc/cpu/utils/nanoflann.hpp

@@ -59,7 +59,7 @@
 #include <vector>
 
 /** Library version: 0xMmP (M=Major,m=minor,P=patch) */
-#define NANOFLANN_VERSION 0x130
+#define NANOFLANN_VERSION 0x132
 
 // Avoid conflicting declaration of min/max macros in windows headers
 #if !defined(NOMINMAX) &&                                                      \
@@ -469,7 +469,8 @@ struct SO2_Adaptor {
   /** Note: this assumes that input angles are already in the range [-pi,pi] */
   template <typename U, typename V>
   inline DistanceType accum_dist(const U a, const V b, const size_t) const {
-    DistanceType result = DistanceType(), PI = pi_const<DistanceType>();
+    DistanceType result = DistanceType();
+    DistanceType PI = pi_const<DistanceType>();
     result = b - a;
     if (result > PI)
       result -= 2 * PI;
@@ -1234,11 +1235,9 @@ public:
     assign(dists, (DIM > 0 ? DIM : BaseClassRef::dim),
            zero); // Fill it with zeros.
     DistanceType distsq = this->computeInitialDistances(*this, vec, dists);
-
     searchLevel(result, vec, BaseClassRef::root_node, distsq, dists,
                 epsError); // "count_leaf" parameter removed since was neither
                            // used nor returned to the user.
-
     return result.full();
   }
 
@@ -1596,11 +1595,9 @@ public:
     assign(dists, (DIM > 0 ? DIM : BaseClassRef::dim),
            static_cast<typename distance_vector_t::value_type>(0));
     DistanceType distsq = this->computeInitialDistances(*this, vec, dists);
-
     searchLevel(result, vec, BaseClassRef::root_node, distsq, dists,
                 epsError); // "count_leaf" parameter removed since was neither
                            // used nor returned to the user.
-
     return result.full();
   }
 
@@ -1936,8 +1933,8 @@ public:
 };
 
 /** An L2-metric KD-tree adaptor for working with data directly stored in an
- * Eigen Matrix, without duplicating the data storage. Each row in the matrix
- * represents a point in the state space.
+ * Eigen Matrix, without duplicating the data storage. You can select whether a
+ * row or column in the matrix represents a point in the state space.
  *
  *  Example of usage:
  * \code
@@ -1951,11 +1948,14 @@ public:
  *  \tparam DIM If set to >0, it specifies a compile-time fixed dimensionality
  * for the points in the data set, allowing more compiler optimizations. \tparam
  * Distance The distance metric to use: nanoflann::metric_L1,
- * nanoflann::metric_L2, nanoflann::metric_L2_Simple, etc.
+ * nanoflann::metric_L2, nanoflann::metric_L2_Simple, etc. \tparam row_major
+ * If set to true the rows of the matrix are used as the points, if set to false
+ * the columns of the matrix are used as the points.
  */
-template <class MatrixType, int DIM = -1, class Distance = nanoflann::metric_L2>
+template <class MatrixType, int DIM = -1, class Distance = nanoflann::metric_L2,
+          bool row_major = true>
 struct KDTreeEigenMatrixAdaptor {
-  typedef KDTreeEigenMatrixAdaptor<MatrixType, DIM, Distance> self_t;
+  typedef KDTreeEigenMatrixAdaptor<MatrixType, DIM, Distance, row_major> self_t;
   typedef typename MatrixType::Scalar num_t;
   typedef typename MatrixType::Index IndexType;
   typedef
@@ -1972,7 +1972,7 @@ struct KDTreeEigenMatrixAdaptor {
                            const std::reference_wrapper<const MatrixType> &mat,
                            const int leaf_max_size = 10)
       : m_data_matrix(mat) {
-    const auto dims = mat.get().cols();
+    const auto dims = row_major ? mat.get().cols() : mat.get().rows();
     if (size_t(dims) != dimensionality)
       throw std::runtime_error(
           "Error: 'dimensionality' must match column count in data matrix");
@@ -2015,12 +2015,18 @@ public:
 
   // Must return the number of data points
   inline size_t kdtree_get_point_count() const {
-    return m_data_matrix.get().rows();
+    if (row_major)
+      return m_data_matrix.get().rows();
+    else
+      return m_data_matrix.get().cols();
   }
 
   // Returns the dim'th component of the idx'th point in the class:
   inline num_t kdtree_get_pt(const IndexType idx, size_t dim) const {
-    return m_data_matrix.get().coeff(idx, IndexType(dim));
+    if (row_major)
+      return m_data_matrix.get().coeff(idx, IndexType(dim));
+    else
+      return m_data_matrix.get().coeff(IndexType(dim), idx);
   }
 
   // Optional bounding-box computation: return false to default to a standard

csrc/cpu/utils/neighbors.cpp

-#include "cloud.h"
-#include "nanoflann.hpp"
-#include <cstdint>
-#include <iostream>
-#include <set>
-#include <thread>
-
-typedef struct thread_struct {
-  void *kd_tree;
-  void *matches;
-  void *queries;
-  size_t *max_count;
-  std::mutex *ct_m;
-  std::mutex *tree_m;
-  size_t start;
-  size_t end;
-  double search_radius;
-  bool small;
-  bool option;
-  size_t k;
-} thread_args;
-
-template <typename scalar_t> void thread_routine(thread_args *targs) {
-  typedef nanoflann::KDTreeSingleIndexAdaptor<
-      nanoflann::L2_Adaptor<scalar_t, PointCloud<scalar_t>>,
-      PointCloud<scalar_t>>
-      my_kd_tree_t;
-  typedef std::vector<std::vector<std::pair<size_t, scalar_t>>> kd_pair;
-  my_kd_tree_t *index = (my_kd_tree_t *)targs->kd_tree;
-  kd_pair *matches = (kd_pair *)targs->matches;
-  PointCloud<scalar_t> *pcd_query = (PointCloud<scalar_t> *)targs->queries;
-  size_t *max_count = targs->max_count;
-  std::mutex *ct_m = targs->ct_m;
-  std::mutex *tree_m = targs->tree_m;
-  double eps;
-  if (targs->small) {
-    eps = 0.000001;
-  } else {
-    eps = 0;
-  }
-  double search_radius = (double)targs->search_radius;
-  size_t start = targs->start;
-  size_t end = targs->end;
-  auto k = targs->k;
-  for (size_t i = start; i < end; i++) {
-
-    std::vector<scalar_t> p0 = *(((*pcd_query).pts)[i]);
-
-    scalar_t *query_pt = new scalar_t[p0.size()];
-    std::copy(p0.begin(), p0.end(), query_pt);
-    (*matches)[i].reserve(*max_count);
-    std::vector<std::pair<size_t, scalar_t>> ret_matches;
-    std::vector<size_t> *knn_ret_matches = new std::vector<size_t>(k);
-    std::vector<scalar_t> *knn_dist_matches = new std::vector<scalar_t>(k);
-
-    tree_m->lock();
-
-    size_t nMatches;
-    if (targs->option) {
-      nMatches = index->radiusSearch(query_pt, (scalar_t)(search_radius + eps),
-                                     ret_matches, nanoflann::SearchParams());
-    } else {
-      nMatches = index->knnSearch(query_pt, k, &(*knn_ret_matches)[0],
-                                  &(*knn_dist_matches)[0]);
-      auto temp = new std::vector<std::pair<size_t, scalar_t>>(
-          (*knn_dist_matches).size());
-      for (size_t j = 0; j < (*knn_ret_matches).size(); j++) {
-        (*temp)[j] =
-            std::make_pair((*knn_ret_matches)[j], (*knn_dist_matches)[j]);
-      }
-      ret_matches = *temp;
-    }
-    tree_m->unlock();
-
-    (*matches)[i] = ret_matches;
-
-    ct_m->lock();
-    if (*max_count < nMatches) {
-      *max_count = nMatches;
-    }
-    ct_m->unlock();
-  }
-}
-
-template <typename scalar_t>
-size_t nanoflann_neighbors(std::vector<scalar_t> &queries,
-                           std::vector<scalar_t> &supports,
-                           std::vector<size_t> &neighbors_indices,
-                           double radius, int dim, int64_t max_num,
-                           int64_t n_threads, int64_t k, int option) {
-
-  const scalar_t search_radius = static_cast<scalar_t>(radius * radius);
-
-  // Counting vector
-  size_t *max_count = new size_t();
-  *max_count = 1;
-
-  size_t ssize = supports.size();
-  // CLoud variable
-  PointCloud<scalar_t> pcd;
-  pcd.set(supports, dim);
-  // Cloud query
-  PointCloud<scalar_t> *pcd_query = new PointCloud<scalar_t>();
-  (*pcd_query).set(queries, dim);
-
-  // Tree parameters
-  nanoflann::KDTreeSingleIndexAdaptorParams tree_params(15 /* max leaf */);
-
-  // KDTree type definition
-  typedef nanoflann::KDTreeSingleIndexAdaptor<
-      nanoflann::L2_Adaptor<scalar_t, PointCloud<scalar_t>>,
-      PointCloud<scalar_t>>
-      my_kd_tree_t;
-  typedef std::vector<std::vector<std::pair<size_t, scalar_t>>> kd_pair;
-
-  // Pointer to trees
-  my_kd_tree_t *index;
-  index = new my_kd_tree_t(dim, pcd, tree_params);
-  index->buildIndex();
-  // Search neigbors indices
-
-  // Search params
-  nanoflann::SearchParams search_params;
-  // search_params.sorted = true;
-  kd_pair *list_matches = new kd_pair((*pcd_query).pts.size());
-
-  // single threaded routine
-  if (n_threads == 1) {
-    size_t i0 = 0;
-    double eps;
-    if (ssize < 10) {
-      eps = 0.000001;
-    } else {
-      eps = 0;
-    }
-
-    for (auto &p : (*pcd_query).pts) {
-      auto p0 = *p;
-      // Find neighbors
-      scalar_t *query_pt = new scalar_t[dim];
-      std::copy(p0.begin(), p0.end(), query_pt);
-
-      (*list_matches)[i0].reserve(*max_count);
-      std::vector<std::pair<size_t, scalar_t>> ret_matches;
-      std::vector<size_t> *knn_ret_matches = new std::vector<size_t>(k);
-      std::vector<scalar_t> *knn_dist_matches = new std::vector<scalar_t>(k);
-
-      size_t nMatches;
-      if (!!(option)) {
-        nMatches =
-            index->radiusSearch(query_pt, (scalar_t)(search_radius + eps),
-                                ret_matches, search_params);
-      } else {
-        nMatches = index->knnSearch(query_pt, (size_t)k, &(*knn_ret_matches)[0],
-                                    &(*knn_dist_matches)[0]);
-        auto temp = new std::vector<std::pair<size_t, scalar_t>>(
-            (*knn_dist_matches).size());
-        for (size_t j = 0; j < (*knn_ret_matches).size(); j++) {
-          (*temp)[j] =
-              std::make_pair((*knn_ret_matches)[j], (*knn_dist_matches)[j]);
-        }
-        ret_matches = *temp;
-      }
-      (*list_matches)[i0] = ret_matches;
-      if (*max_count < nMatches)
-        *max_count = nMatches;
-      i0++;
-    }
-  } else { // Multi-threaded routine
-    std::mutex *mtx = new std::mutex();
-    std::mutex *mtx_tree = new std::mutex();
-
-    size_t n_queries = (*pcd_query).pts.size();
-    size_t actual_threads =
-        std::min((long long)n_threads, (long long)n_queries);
-
-    std::vector<std::thread *> tid(actual_threads);
-
-    size_t start, end;
-    size_t length;
-    if (n_queries) {
-      length = 1;
-    } else {
-      auto res = std::lldiv((long long)n_queries, (long long)n_threads);
-      length = (size_t)res.quot;
-    }
-    for (size_t t = 0; t < actual_threads; t++) {
-      start = t * length;
-      if (t == actual_threads - 1) {
-        end = n_queries;
-      } else {
-        end = (t + 1) * length;
-      }
-      thread_args *targs = new thread_args();
-      targs->kd_tree = index;
-      targs->matches = list_matches;
-      targs->max_count = max_count;
-      targs->ct_m = mtx;
-      targs->tree_m = mtx_tree;
-      targs->search_radius = search_radius;
-      targs->queries = pcd_query;
-      targs->start = start;
-      targs->end = end;
-      if (ssize < 10) {
-        targs->small = true;
-      } else {
-        targs->small = false;
-      }
-      targs->option = !!(option);
-      targs->k = k;
-      std::thread *temp = new std::thread(thread_routine<scalar_t>, targs);
-      tid[t] = temp;
-    }
-
-    for (size_t t = 0; t < actual_threads; t++) {
-      tid[t]->join();
-    }
-  }
-
-  // Reserve the memory
-  if (max_num > 0) {
-    *max_count = max_num;
-  }
-
-  size_t size = 0; // total number of edges
-  for (auto &inds : *list_matches) {
-    if (inds.size() <= *max_count)
-      size += inds.size();
-    else
-      size += *max_count;
-  }
-
-  neighbors_indices.resize(size * 2);
-  size_t i1 = 0; // index of the query points
-  size_t u = 0;  // curent index of the neighbors_indices
-  for (auto &inds : *list_matches) {
-    for (size_t j = 0; j < *max_count; j++) {
-      if (j < inds.size()) {
-        neighbors_indices[u] = inds[j].first;
-        neighbors_indices[u + 1] = i1;
-        u += 2;
-      }
-    }
-    i1++;
-  }
-
-  return *max_count;
-}
-
-template <typename scalar_t>
-size_t batch_nanoflann_neighbors(std::vector<scalar_t> &queries,
-                                 std::vector<scalar_t> &supports,
-                                 std::vector<long> &q_batches,
-                                 std::vector<long> &s_batches,
-                                 std::vector<size_t> &neighbors_indices,
-                                 double radius, int dim, int64_t max_num,
-                                 int64_t k, int option) {
-
-  // Indices.
-  size_t i0 = 0;
-
-  // Square radius.
-  const scalar_t r2 = static_cast<scalar_t>(radius * radius);
-
-  // Counting vector.
-  size_t max_count = 0;
-
-  // Batch index.
-  size_t b = 0;
-  size_t sum_qb = 0;
-  size_t sum_sb = 0;
-
-  double eps;
-  if (supports.size() < 10) {
-    eps = 0.000001;
-  } else {
-    eps = 0;
-  }
-  // Nanoflann related variables.
-
-  // Cloud variable.
-  PointCloud<scalar_t> current_cloud;
-  PointCloud<scalar_t> query_pcd;
-  query_pcd.set(queries, dim);
-  std::vector<std::vector<std::pair<size_t, scalar_t>>> all_inds_dists(
-      query_pcd.pts.size());
-
-  // Tree parameters.
-  nanoflann::KDTreeSingleIndexAdaptorParams tree_params(10 /* max leaf */);
-
-  // KDTree type definition.
-  typedef nanoflann::KDTreeSingleIndexAdaptor<
-      nanoflann::L2_Adaptor<scalar_t, PointCloud<scalar_t>>,
-      PointCloud<scalar_t>>
-      my_kd_tree_t;
-
-  // Pointer to trees.
-  my_kd_tree_t *index;
-  // Build KDTree for the first batch element.
-  current_cloud.set_batch(supports, sum_sb, s_batches[b], dim);
-  index = new my_kd_tree_t(dim, current_cloud, tree_params);
-  index->buildIndex();
-  // Search neigbors indices.
-  // Search params.
-  nanoflann::SearchParams search_params;
-  search_params.sorted = true;
-
-  for (auto &p : query_pcd.pts) {
-    auto p0 = *p;
-    // Check if we changed batch.
-
-    scalar_t *query_pt = new scalar_t[dim];
-    std::copy(p0.begin(), p0.end(), query_pt);
-
-    if (i0 == sum_qb + q_batches[b]) {
-      sum_qb += q_batches[b];
-      sum_sb += s_batches[b];
-      b++;
-
-      // Change the points.
-      current_cloud.pts.clear();
-      current_cloud.set_batch(supports, sum_sb, s_batches[b], dim);
-      // Build KDTree of the current element of the batch.
-      delete index;
-      index = new my_kd_tree_t(dim, current_cloud, tree_params);
-      index->buildIndex();
-    }
-    // Initial guess of neighbors size.
-    all_inds_dists[i0].reserve(max_count);
-
-    // Find neighbors.
-    size_t nMatches;
-    if (!!option) {
-      nMatches = index->radiusSearch(query_pt, r2 + eps, all_inds_dists[i0],
-                                     search_params);
-      // Update max count.
-    } else {
-      std::vector<size_t> *knn_ret_matches = new std::vector<size_t>(k);
-      std::vector<scalar_t> *knn_dist_matches = new std::vector<scalar_t>(k);
-      nMatches = index->knnSearch(query_pt, (size_t)k, &(*knn_ret_matches)[0],
-                                  &(*knn_dist_matches)[0]);
-      auto temp = new std::vector<std::pair<size_t, scalar_t>>(
-          (*knn_dist_matches).size());
-      for (size_t j = 0; j < (*knn_ret_matches).size(); j++) {
-        (*temp)[j] =
-            std::make_pair((*knn_ret_matches)[j], (*knn_dist_matches)[j]);
-      }
-      all_inds_dists[i0] = *temp;
-    }
-    if (nMatches > max_count)
-      max_count = nMatches;
-    i0++;
-  }
-
-  // How many neighbors do we keep.
-  if (max_num > 0) {
-    max_count = max_num;
-  }
-
-  size_t size = 0; // Total number of edges.
-  for (auto &inds_dists : all_inds_dists) {
-    if (inds_dists.size() <= max_count)
-      size += inds_dists.size();
-    else
-      size += max_count;
-  }
-
-  neighbors_indices.resize(size * 2);
-  i0 = 0;
-  sum_sb = 0;
-  sum_qb = 0;
-  b = 0;
-  size_t u = 0;
-  for (auto &inds_dists : all_inds_dists) {
-    if (i0 == sum_qb + q_batches[b]) {
-      sum_qb += q_batches[b];
-      sum_sb += s_batches[b];
-      b++;
-    }
-    for (size_t j = 0; j < max_count; j++) {
-      if (j < inds_dists.size()) {
-        neighbors_indices[u] = inds_dists[j].first + sum_sb;
-        neighbors_indices[u + 1] = i0;
-        u += 2;
-      }
-    }
-    i0++;
-  }
-
-  return max_count;
-}