text: /*
   * Copyright (c) 2020 NVIDIA CORPORATION.
   * Copyright (c) Chris Choy (chrischoy@ai.stanford.edu).
   *
   * Permission is hereby granted, free of charge, to any person obtaining a copy
   * of this software and associated documentation files (the "Software"), to deal
   * in the Software without restriction, including without limitation the rights
   * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
   * copies of the Software, and to permit persons to whom the Software is
   * furnished to do so, subject to the following conditions:
   *
   * The above copyright notice and this permission notice shall be included in
   * all copies or substantial portions of the Software.
   *
   * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
   * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
   * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
   * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
   * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
   * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
   * IN THE SOFTWARE.
   *
   * Please cite "4D Spatio-Temporal ConvNets: Minkowski Convolutional Neural
   * Networks", CVPR'19 (https://arxiv.org/abs/1904.08755) if you use any part
   * of the code.
   */
  #include "coordinate_map_manager.hpp"
  #include "coordinate_map_key.hpp"
  #include "errors.hpp"
  #include "kernel_region.hpp"
  #include "utils.hpp"
  
  #include <pybind11/pybind11.h>
  #include <string>
  #include <unordered_map>
  
  namespace py = pybind11;
  
  namespace minkowski {
  
  namespace detail {
  
  default_types::stride_type zeros(size_t const len) { return _fill_vec<0>(len); }
  
  default_types::stride_type ones(size_t const len) { return _fill_vec<1>(len); }
  
  } // namespace detail
  
  /*******************************
   * Initialization
   *******************************/
  
  namespace detail {
  
  template <typename coordinate_type, typename coordinate_field_type>
  struct insert_and_map_functor<coordinate_type, coordinate_field_type,
                                std::allocator, CoordinateMapCPU> {
  
    std::pair<at::Tensor, at::Tensor>
    operator()(coordinate_map_key_type &map_key, at::Tensor const &th_coordinate,
               CoordinateMapManager<coordinate_type, coordinate_field_type,
                                    std::allocator, CoordinateMapCPU> &manager) {
      LOG_DEBUG("initialize_and_map");
      uint32_t const N = th_coordinate.size(0);
      uint32_t const coordinate_size = th_coordinate.size(1);
      coordinate_type *p_coordinate = th_coordinate.data_ptr<coordinate_type>();
      auto map = CoordinateMapCPU<coordinate_type, std::allocator>(
          N, coordinate_size, map_key.first);
      auto map_inverse_map = map.template insert_and_map<true>(
          p_coordinate, p_coordinate + N * coordinate_size);
      LOG_DEBUG("mapping size:", map_inverse_map.first.size());
  
      // insert moves map
      THRUST_CHECK(manager.insert(map_key, map));
  
      auto const &mapping = map_inverse_map.first;
      auto const &inverse_mapping = map_inverse_map.second;
  
      // return tensors
      at::Tensor th_mapping = torch::empty(
          {(int64_t)mapping.size()},
          torch::TensorOptions().requires_grad(false).dtype(torch::kInt64));
      at::Tensor th_inverse_mapping = torch::empty(
          {(int64_t)inverse_mapping.size()},
          torch::TensorOptions().requires_grad(false).dtype(torch::kInt64));
  
      // copy_n to int to long
      int64_t *p_mapping = th_mapping.data_ptr<int64_t>();
      for (default_types::index_type i = 0; i < mapping.size(); ++i) {
        p_mapping[i] = mapping[i];
      }
  
      int64_t *p_inverse_mapping = th_inverse_mapping.data_ptr<int64_t>();
      for (default_types::index_type i = 0; i < inverse_mapping.size(); ++i) {
        p_inverse_mapping[i] = inverse_mapping[i];
      }
  
      return std::make_pair(std::move(th_mapping), std::move(th_inverse_mapping));
    }
  };
  
  template <typename coordinate_type, typename coordinate_field_type>
  struct insert_field_functor<
      coordinate_type, coordinate_field_type, std::allocator, CoordinateMapCPU,
      CoordinateFieldMapCPU<coordinate_field_type, coordinate_type,
                            std::allocator>> {
  
    void
    operator()(coordinate_map_key_type &map_key, at::Tensor const &th_coordinate,
               CoordinateMapManager<coordinate_type, coordinate_field_type,
                                    std::allocator, CoordinateMapCPU> &manager) {
      LOG_DEBUG("insert field");
      uint32_t const N = th_coordinate.size(0);
      uint32_t const coordinate_size = th_coordinate.size(1);
      coordinate_field_type *p_coordinate =
          th_coordinate.data_ptr<coordinate_field_type>();
      auto map = CoordinateFieldMapCPU<coordinate_field_type, coordinate_type,
                                       std::allocator>(N, coordinate_size,
                                                       map_key.first);
      THRUST_CHECK(map.insert(p_coordinate, p_coordinate + N * coordinate_size));
  
      LOG_DEBUG("insert map with tensor_stride", map_key.first);
      manager.insert_field_map(map_key, map);
    }
  };
  
  } // namespace detail
  
  /*
   * coords: coordinates in IntTensor
   * mapping: output mapping in IntTensor
   * tensor_strides: current tensor strides this coords will be initializeds
   * force_creation: even when there's a duplicate coords with the same tensor
   *                 strides.
   * force_remap: if there's duplicate coords, remap
   * allow_duplicate_coords: create map when there are duplicates in the
   * coordinates
   */
  template <typename coordinate_type, typename coordinate_field_type,
            template <typename C> class TemplatedAllocator,
            template <typename T, template <typename Q> class A>
            class CoordinateMapType>
  py::object CoordinateMapManager<coordinate_type, coordinate_field_type,
                                  TemplatedAllocator, CoordinateMapType>::
      insert_field(at::Tensor const &coordinates,
                   default_types::stride_type const tensor_stride,
                   std::string const string_id) {
  
    torch::TensorArg arg_coordinate(coordinates, "coordinates", 0);
    torch::CheckedFrom c = "initialize";
    torch::checkContiguous(c, arg_coordinate);
  
    // must match coordinate_type
    torch::checkScalarType(c, arg_coordinate, torch::kFloat);
    torch::checkBackend(c, arg_coordinate.tensor,
                        detail::is_cpu_coordinate_map<CoordinateMapType>::value
                            ? torch::Backend::CPU
                            : torch::Backend::CUDA);
    torch::checkDim(c, arg_coordinate, 2);
  
    auto const coordinate_size = (index_type)coordinates.size(1);
  
    // Basic assertions
    ASSERT(coordinate_size - 1 == tensor_stride.size(),
           "The coordinate dimension (coordinate_size - 1):", coordinate_size - 1,
           " must match the size of tensor stride:", ArrToString(tensor_stride));
  
    // generate the map_key
    coordinate_map_key_type map_key = std::make_pair(tensor_stride, string_id);
    if (m_field_coordinates.find(map_key) != m_field_coordinates.end()) {
      LOG_DEBUG("CoordinateMapKey collision detected:", map_key,
                "generating new string id.");
      map_key = get_random_string_id(tensor_stride, string_id);
    }
  
    LOG_DEBUG("initializing a map with tensor stride:", map_key.first,
              "string id:", map_key.second);
    // Create the concurrent coords map
    detail::insert_field_functor<coordinate_type, coordinate_field_type,
                                 TemplatedAllocator, CoordinateMapType,
                                 field_map_type>()(map_key, coordinates, *this);
  
    py::object py_key = py::cast(new CoordinateMapKey(coordinate_size, map_key));
  
    return py_key;
  }
  
  /* to_sparse_and_map */
  /*
   * coords: coordinates in IntTensor
   * tensor_strides: current tensor strides this coords will be initializeds
   */
  template <typename coordinate_type, typename coordinate_field_type,
            template <typename C> class TemplatedAllocator,
            template <typename T, template <typename Q> class A>
            class CoordinateMapType>
  std::pair<py::object, std::pair<at::Tensor, at::Tensor>>
  CoordinateMapManager<coordinate_type, coordinate_field_type, TemplatedAllocator,
                       CoordinateMapType>::
      field_to_sparse_insert_and_map(
          CoordinateMapKey const *p_in_field_map_key,
          default_types::stride_type const sparse_tensor_stride,
          std::string const sparse_tensor_string_id) {
    auto const coordinate_size = p_in_field_map_key->get_coordinate_size();
    // Basic assertions
    ASSERT(coordinate_size - 1 == sparse_tensor_stride.size(),
           "The coordinate dimension (coordinate_size - 1):", coordinate_size - 1,
           " must match the size of tensor stride:",
           ArrToString(sparse_tensor_stride));
  
    // Find coordinate field
    auto const it = m_field_coordinates.find(p_in_field_map_key->get_key());
    ASSERT(it != m_field_coordinates.end(), ERROR_MAP_NOT_FOUND);
    auto const &field_map = it->second;
  
    auto options = torch::TensorOptions().dtype(torch::kInt).requires_grad(false);
  
    if (!detail::is_cpu_coordinate_map<CoordinateMapType>::value) {
  #ifndef CPU_ONLY
      auto device_id = at::cuda::current_device();
      options = options.device(torch::kCUDA, device_id);
  #else
      ASSERT(false, ERROR_CPU_ONLY);
  #endif
    }
  
    // generate the map_key
    coordinate_map_key_type map_key =
        std::make_pair(sparse_tensor_stride, sparse_tensor_string_id);
    if (m_coordinate_maps.find(map_key) != m_coordinate_maps.end()) {
      LOG_DEBUG("CoordinateMapKey collision detected:", map_key,
                "generating new string id.");
      map_key =
          get_random_string_id(sparse_tensor_stride, sparse_tensor_string_id);
    }
  
    LOG_DEBUG("initializing a field map with tensor stride:", map_key.first,
              "string id:", map_key.second);
  
    // Quantize the field with tensor stride.
    // The coordinate must be a tensor. Wrap a pointer with a tensor.
    at::Tensor int_coordinates =
        at::empty({field_map.size(), coordinate_size}, options);
    field_map.quantize_coordinates(int_coordinates.data_ptr<coordinate_type>(),
                                   sparse_tensor_stride);
  
    auto const map_inverse_map =
        detail::insert_and_map_functor<coordinate_type, coordinate_field_type,
                                       TemplatedAllocator, CoordinateMapType>()(
            map_key, int_coordinates, *this);
  
    auto const field_to_sparse_map_key =
        std::pair<coordinate_map_key_type, coordinate_map_key_type>{
            p_in_field_map_key->get_key(), map_key};
  
    auto result = m_field_to_sparse_maps.insert(
        std::pair<
            const std::pair<coordinate_map_key_type, coordinate_map_key_type>,
            const std::pair<at::Tensor, at::Tensor>>{field_to_sparse_map_key,
                                                     map_inverse_map});
    LOG_DEBUG("field to sparse tensor map insertion", result.second);
  
    py::object py_key = py::cast(new CoordinateMapKey(coordinate_size, map_key));
  
    return std::make_pair(py_key, map_inverse_map);
  }
  
  /*
   * Return existing field to sparse map. Raise an except if the map doesn't
   * exist
   */
  template <typename coordinate_type, typename coordinate_field_type,
            template <typename C> class TemplatedAllocator,
            template <typename T, template <typename Q> class A>
            class CoordinateMapType>
  std::pair<at::Tensor, at::Tensor>
  CoordinateMapManager<coordinate_type, coordinate_field_type, TemplatedAllocator,
                       CoordinateMapType>::
      get_field_to_sparse_map(CoordinateMapKey const *p_field_key,
                              CoordinateMapKey const *p_sparse_key) const {
    auto key = std::pair<coordinate_map_key_type, coordinate_map_key_type>{
        p_field_key->get_key(), p_sparse_key->get_key()};
    auto it = m_field_to_sparse_maps.find(key);
    ASSERT(it != m_field_to_sparse_maps.end(),
           "Field To Sparse Map doesn't exist");
    return it->second;
  }
  
  /*
   * Create a field to sparse map if it doesn't exist.
   */
  template <typename coordinate_type, typename coordinate_field_type,
            template <typename C> class TemplatedAllocator,
            template <typename T, template <typename Q> class A>
            class CoordinateMapType>
  std::pair<at::Tensor, at::Tensor>
  CoordinateMapManager<coordinate_type, coordinate_field_type, TemplatedAllocator,
                       CoordinateMapType>::
      field_to_sparse_map(CoordinateMapKey const *p_in_field_map_key,
                          CoordinateMapKey const *p_out_sparse_map_key) {
  
    auto const coordinate_size = p_in_field_map_key->get_coordinate_size();
    // Basic assertions
    ASSERT(coordinate_size == p_out_sparse_map_key->get_coordinate_size(),
           "The coordinate dimension mismatch.", coordinate_size,
           "!=", p_out_sparse_map_key->get_coordinate_size());
  
    // Find coordinate field
    auto const it_field = m_field_coordinates.find(p_in_field_map_key->get_key());
    ASSERT(it_field != m_field_coordinates.end(), ERROR_MAP_NOT_FOUND);
    auto const &field_map = it_field->second;
    auto const it_sparse =
        m_coordinate_maps.find(p_out_sparse_map_key->get_key());
    ASSERT(it_sparse != m_coordinate_maps.end(), ERROR_MAP_NOT_FOUND);
    auto const &sparse_map = it_sparse->second;
  
    auto options = torch::TensorOptions().dtype(torch::kInt).requires_grad(false);
  
    if (!detail::is_cpu_coordinate_map<CoordinateMapType>::value) {
  #ifndef CPU_ONLY
      auto device_id = at::cuda::current_device();
      options = options.device(torch::kCUDA, device_id);
  #else
      ASSERT(false, ERROR_CPU_ONLY);
  #endif
    }
  
    auto const map_inverse_map =
        sparse_map.field_map(field_map.const_coordinate_data(), field_map.size());
  
    auto const field_to_sparse_map_key =
        std::pair<coordinate_map_key_type, coordinate_map_key_type>{
            p_in_field_map_key->get_key(), p_out_sparse_map_key->get_key()};
  
    auto result = m_field_to_sparse_maps.insert(
        std::pair<
            const std::pair<coordinate_map_key_type, coordinate_map_key_type>,
            const std::pair<at::Tensor, at::Tensor>>{field_to_sparse_map_key,
                                                     map_inverse_map});
    LOG_DEBUG("field to sparse tensor map insertion", result.second);
  
    return map_inverse_map;
  }
  
  /*
   * coords: coordinates in IntTensor
   * tensor_strides: current tensor strides this coords will be initializeds
   */
  template <typename coordinate_type, typename coordinate_field_type,
            template <typename C> class TemplatedAllocator,
            template <typename T, template <typename Q> class A>
            class CoordinateMapType>
  std::pair<py::object, std::pair<at::Tensor, at::Tensor>>
  CoordinateMapManager<coordinate_type, coordinate_field_type, TemplatedAllocator,
                       CoordinateMapType>::
      insert_and_map(at::Tensor const &coordinate,
                     default_types::stride_type const tensor_stride,
                     std::string const string_id) {
  
    torch::TensorArg arg_coordinate(coordinate, "coordinates", 0);
    torch::CheckedFrom c = "initialize";
    torch::checkContiguous(c, arg_coordinate);
    // must match coordinate_type
    torch::checkScalarType(c, arg_coordinate, torch::kInt);
    torch::checkBackend(c, arg_coordinate.tensor,
                        detail::is_cpu_coordinate_map<CoordinateMapType>::value
                            ? torch::Backend::CPU
                            : torch::Backend::CUDA);
    torch::checkDim(c, arg_coordinate, 2);
  
    auto const coordinate_size = (index_type)coordinate.size(1);
  
    // Basic assertions
    ASSERT(coordinate_size - 1 == tensor_stride.size(),
           "The coordinate dimension (coordinate_size - 1):", coordinate_size - 1,
           " must match the size of tensor stride:", ArrToString(tensor_stride));
  
    // generate the map_key
    coordinate_map_key_type map_key = std::make_pair(tensor_stride, string_id);
    if (m_coordinate_maps.find(map_key) != m_coordinate_maps.end()) {
      LOG_DEBUG("CoordinateMapKey collision detected:", map_key,
                "generating new string id.");
      map_key = get_random_string_id(tensor_stride, string_id);
    }
  
    LOG_DEBUG("initializing a map with tensor stride:", map_key.first,
              "string id:", map_key.second);
    // Create the concurrent coords map
    auto const map_inverse_map =
        detail::insert_and_map_functor<coordinate_type, coordinate_field_type,
                                       TemplatedAllocator, CoordinateMapType>()(
            map_key, coordinate, *this);
  
    LOG_DEBUG("map_inverse_map initialized");
    py::object py_key = py::cast(new CoordinateMapKey(coordinate_size, map_key));
    LOG_DEBUG("py key initialized");
  
    return std::make_pair(py_key, map_inverse_map);
  }
  
  // stride
  template <typename coordinate_type, typename coordinate_field_type,
            template <typename C> class TemplatedAllocator,
            template <typename T, template <typename Q> class A>
            class CoordinateMapType>
  std::pair<coordinate_map_key_type, bool> CoordinateMapManager<
      coordinate_type, coordinate_field_type, TemplatedAllocator,
      CoordinateMapType>::stride(coordinate_map_key_type const &in_map_key,
                                 stride_type const &kernel_stride,
                                 std::string const string_id) {
    ASSERT(exists(in_map_key), ERROR_MAP_NOT_FOUND);
    // check if the key exists.
    LOG_DEBUG("In tensor stride:", in_map_key.first,
              "kernel stride:", kernel_stride);
    coordinate_map_key_type out_map_key(
        detail::stride_tensor_stride(in_map_key.first, kernel_stride, false),
        string_id == "" ? in_map_key.second : string_id);
    LOG_DEBUG("Out stride map key:", out_map_key);
    bool const exists_out_map = exists(out_map_key);
    if (!exists_out_map) {
      // operator[] required mapped_type(), which is not defined.
      // ASSERTION already checked that in_map_key exists.
      map_type const &in_map = m_coordinate_maps.find(in_map_key)->second;
      map_type out_map = in_map.stride(kernel_stride);
      insert(out_map_key, out_map);
    }
    // (key, new map generated flag)
    return std::make_pair(out_map_key, !exists_out_map);
  }
  
  template <typename coordinate_type, typename coordinate_field_type,
            template <typename C> class TemplatedAllocator,
            template <typename T, template <typename Q> class A>
            class CoordinateMapType>
  std::pair<coordinate_map_key_type, bool>
  CoordinateMapManager<coordinate_type, coordinate_field_type, TemplatedAllocator,
                       CoordinateMapType>::
      stride_region(coordinate_map_key_type const &in_map_key,
                    cpu_kernel_region<coordinate_type> &kernel,
                    stride_type const &out_tensor_stride,
                    bool const expand_coordinates) {
    ASSERT(exists(in_map_key), ERROR_MAP_NOT_FOUND);
    LOG_DEBUG("stride_region");
    // kernel.tensor_stride must be set to out tensor stride.
    // stride_type out_tensor_stride{kernel.tensor_stride(),
    //                               kernel.tensor_stride() +
    //                                   kernel.coordinate_size() - 1};
  
    // check if the key exists.
    coordinate_map_key_type out_map_key(out_tensor_stride, "");
    bool const exists_out_map = exists(out_map_key);
    if (!exists_out_map || expand_coordinates) {
      LOG_DEBUG("Create a new stride region map for tensor_stride:",
                out_tensor_stride);
      map_type const &in_map = m_coordinate_maps.find(in_map_key)->second;
      map_type out_map = in_map.stride_region(kernel, out_tensor_stride);
      if (exists_out_map) {
        LOG_DEBUG("coordinate map exists for tensor_stride:", out_tensor_stride);
        out_map_key = get_random_string_id(out_tensor_stride, "");
        LOG_DEBUG("created a random key:", out_map_key);
      }
      insert(out_map_key, out_map);
    }
    // (key, new map generated flag)
    return std::make_pair(out_map_key, !exists_out_map || expand_coordinates);
  }
  
  template <typename coordinate_type, typename coordinate_field_type,
            template <typename C> class TemplatedAllocator,
            template <typename T, template <typename Q> class A>
            class CoordinateMapType>
  std::pair<coordinate_map_key_type, bool>
  CoordinateMapManager<coordinate_type, coordinate_field_type, TemplatedAllocator,
                       CoordinateMapType>::origin() {
    ASSERT(m_coordinate_maps.size() > 0, "No coordinate map found");
    // check if the key exists.
    map_type const &random_map = m_coordinate_maps.begin()->second;
    stride_type origin_tensor_stride(random_map.coordinate_size() - 1);
    std::for_each(origin_tensor_stride.begin(), origin_tensor_stride.end(),
                  [](auto &i) { i = 0; });
    LOG_DEBUG("origin tensor stride:", origin_tensor_stride);
  
    coordinate_map_key_type origin_map_key(origin_tensor_stride, "");
    bool const exists_origin_map = exists(origin_map_key);
    if (!exists_origin_map) {
      LOG_DEBUG("origin coordinate map not found");
      map_type const *p_min_coordinate_map{nullptr};
      size_type min_size = std::numeric_limits<size_type>::max();
      for (auto map_it = m_coordinate_maps.begin();
           map_it != m_coordinate_maps.end(); ++map_it) {
        if (min_size > map_it->second.size()) {
          p_min_coordinate_map = &(map_it->second);
        }
      }
  
      if (p_min_coordinate_map != nullptr) {
        map_type origin_map = p_min_coordinate_map->origin();
        LOG_DEBUG("origin map with size:", origin_map.size(), " inserted");
        insert(origin_map_key, origin_map);
      } else {
        ASSERT(false, "Invalid origin map");
      }
    }
  
    LOG_DEBUG("return origin()");
    // (key, new map generated flag)
    return std::make_pair(origin_map_key, !exists_origin_map);
  }
  
  template <typename coordinate_type, typename coordinate_field_type,
            template <typename C> class TemplatedAllocator,
            template <typename T, template <typename Q> class A>
            class CoordinateMapType>
  std::pair<coordinate_map_key_type, bool>
  CoordinateMapManager<coordinate_type, coordinate_field_type, TemplatedAllocator,
                       CoordinateMapType>::origin_field() {
    ASSERT(m_field_coordinates.size() > 0, "No coordinate map found");
    // check if the key exists.
    field_map_type const &random_map = m_field_coordinates.begin()->second;
    stride_type origin_tensor_stride(random_map.coordinate_size() - 1);
    std::for_each(origin_tensor_stride.begin(), origin_tensor_stride.end(),
                  [](auto &i) { i = 0; });
    LOG_DEBUG("origin tensor stride:", origin_tensor_stride);
  
    coordinate_map_key_type origin_map_key(origin_tensor_stride, "");
    bool const exists_origin_map = exists(origin_map_key);
  
    if (!exists_origin_map) {
      LOG_DEBUG("origin coordinate map not found");
      field_map_type const *p_min_coordinate_map{nullptr};
      size_type min_size = std::numeric_limits<size_type>::max();
      for (auto map_it = m_field_coordinates.begin();
           map_it != m_field_coordinates.end(); ++map_it) {
        if (min_size > map_it->second.size()) {
          p_min_coordinate_map = &(map_it->second);
        }
      }
  
      if (p_min_coordinate_map != nullptr) {
        map_type origin_map = p_min_coordinate_map->origin();
        LOG_DEBUG("origin map with size:", origin_map.size(), " inserted");
        insert(origin_map_key, origin_map);
      } else {
        ASSERT(false, "Invalid origin map");
      }
    }
  
    // (key, new map generated flag)
    return std::make_pair(origin_map_key, !exists_origin_map);
  }
  
  template <typename coordinate_type, typename coordinate_field_type,
            template <typename C> class TemplatedAllocator,
            template <typename T, template <typename Q> class A>
            class CoordinateMapType>
  coordinate_map_key_type
  CoordinateMapManager<coordinate_type, coordinate_field_type, TemplatedAllocator,
                       CoordinateMapType>::prune(coordinate_map_key_type const
                                                     &in_key,
                                                 bool const *keep_begin,
                                                 bool const *keep_end) {
    auto const map_it = m_coordinate_maps.find(in_key);
    ASSERT(map_it != m_coordinate_maps.end(), ERROR_MAP_NOT_FOUND);
  
    // create a coordinate_map_key
    coordinate_map_key_type map_key = std::make_pair(in_key.first, "pruned");
    if (m_coordinate_maps.find(map_key) != m_coordinate_maps.end()) {
      map_key = get_random_string_id(map_key.first, map_key.second);
    }
  
    map_type pruned_map = map_it->second.prune(keep_begin, keep_end);
    LOG_DEBUG("pruned map with size:", pruned_map.size(), " inserted");
    insert(map_key, pruned_map);
  
    return map_key;
  }
  
  // Kernel map
  
  namespace detail {
  
  template <typename coordinate_type>
  struct kernel_map_functor<coordinate_type, std::allocator, CoordinateMapCPU,
                            cpu_kernel_map> {
  
    cpu_kernel_map
    operator()(CoordinateMapCPU<coordinate_type, std::allocator> const &in_map,
               CoordinateMapCPU<coordinate_type, std::allocator> const &out_map,
               CUDAKernelMapMode::Mode kernel_map_mode,
               cpu_kernel_region<coordinate_type> &kernel) {
      return in_map.kernel_map(out_map, kernel);
    }
  };
  
  template <typename coordinate_type>
  struct stride_map_functor<coordinate_type, std::allocator, CoordinateMapCPU,
                            cpu_kernel_map> {
  
    cpu_kernel_map
    operator()(CoordinateMapCPU<coordinate_type, std::allocator> const &in_map,
               CoordinateMapCPU<coordinate_type, std::allocator> const &out_map,
               default_types::stride_type const &out_tensor_stride) {
      return in_map.stride_map(out_map, out_tensor_stride);
    }
  };
  
  // a partial specialization functor for kernel map in/out swap
  template <> struct swap_in_out_map_functor<cpu_kernel_map> {
  
    cpu_kernel_map operator()(cpu_kernel_map const &kernel_map) {
      return std::make_pair(kernel_map.second, kernel_map.first);
    }
  };
  
  template <typename coordinate_type>
  struct empty_map_functor<coordinate_type, std::allocator, CoordinateMapCPU,
                           cpu_kernel_map> {
  
    cpu_kernel_map operator()() { return cpu_kernel_map{}; }
  };
  
  } // namespace detail
  
  /*
   * Given tensor_stride_src and tensor_stride_dst, find the respective coord_maps
   * and return the indices of the coord_map_ind in coord_map_dst
   */
  template <typename coordinate_type, typename coordinate_field_type,
            template <typename C> class TemplatedAllocator,
            template <typename T, template <typename Q> class A>
            class CoordinateMapType>
  typename CoordinateMapManager<coordinate_type, coordinate_field_type,
                                TemplatedAllocator,
                                CoordinateMapType>::kernel_map_type const &
  CoordinateMapManager<
      coordinate_type, coordinate_field_type, TemplatedAllocator,
      CoordinateMapType>::kernel_map(CoordinateMapKey const *p_in_map_key,
                                     CoordinateMapKey const *p_out_map_key) {
    // when kernel has volume 1
    auto const &map_it = m_coordinate_maps.find(p_in_map_key->get_key());
    ASSERT(map_it != m_coordinate_maps.end(), ERROR_MAP_NOT_FOUND);
    auto const coordinate_size = map_it->second.coordinate_size();
    auto const one_vec = detail::ones(coordinate_size - 1);
    auto const offset = torch::empty(
        {0}, torch::TensorOptions().dtype(torch::kInt32).requires_grad(false));
  
    return kernel_map(p_in_map_key, p_out_map_key, one_vec, one_vec, one_vec,
                      RegionType::HYPER_CUBE, offset, false, false);
  }
  
  /*
   * Given tensor_stride_src and tensor_stride_dst, find the respective coord_maps
   * and return the indices of the coord_map_ind in coord_map_dst
   */
  template <typename coordinate_type, typename coordinate_field_type,
            template <typename C> class TemplatedAllocator,
            template <typename T, template <typename Q> class A>
            class CoordinateMapType>
  typename CoordinateMapManager<coordinate_type, coordinate_field_type,
                                TemplatedAllocator,
                                CoordinateMapType>::kernel_map_type const &
  CoordinateMapManager<
      coordinate_type, coordinate_field_type, TemplatedAllocator,
      CoordinateMapType>::kernel_map(CoordinateMapKey const *p_in_map_key,
                                     CoordinateMapKey const *p_out_map_key,
                                     stride_type const &kernel_size, //
                                     stride_type const &kernel_stride,
                                     stride_type const &kernel_dilation,
                                     RegionType::Type const region_type,
                                     at::Tensor const &offset, bool is_transpose,
                                     bool is_pool) {
    ASSERT(region_type != RegionType::CUSTOM, "Not implemented yet.");
    if (region_type == RegionType::CUSTOM)
      ASSERT(offset.is_cuda() ==
                 !detail::is_cpu_coordinate_map<CoordinateMapType>::value,
             "Invalid device for offset");
  
    size_type kernel_dim = kernel_size.size();
  
    ASSERT(kernel_dim == kernel_stride.size(), "kernel size mismatch");
    ASSERT(kernel_dim == kernel_dilation.size(), "kernel size mismatch");
  
    // in_coords_key->tensor_stride * kernel_stride ==
    // out_coords_key->tensor_stride
  
    kernel_map_key_type const kernel_map_key =
        std::make_tuple(p_in_map_key->get_key(), p_out_map_key->get_key(), // maps
                        kernel_size, kernel_stride, kernel_dilation, // kernels
                        region_type, is_transpose, is_pool);
  
    const auto &kernel_map_iter = m_kernel_maps.find(kernel_map_key);
    LOG_DEBUG("set kernel map key for kernel map:", p_in_map_key->get_key(), "->",
              p_out_map_key->get_key());
  
    if (kernel_map_iter == m_kernel_maps.end()) {
      // create a kernel map if it exists
      auto const in_map_it = m_coordinate_maps.find(p_in_map_key->get_key());
      auto const out_map_it = m_coordinate_maps.find(p_out_map_key->get_key());
  
      ASSERT(in_map_it != m_coordinate_maps.end(), "in_map", ERROR_MAP_NOT_FOUND);
      ASSERT(out_map_it != m_coordinate_maps.end(), "out_map",
             ERROR_MAP_NOT_FOUND);
  
      auto const &in_map = in_map_it->second;
      auto const &out_map = out_map_it->second;
  
      LOG_DEBUG("coordinate_size:", in_map.coordinate_size(),
                "in tensor_stride:", in_map.get_tensor_stride(),
                "out tensor_stride:", out_map.get_tensor_stride());
  
      // +1 for batch index
      ASSERT(kernel_dim + 1 == in_map.coordinate_size(), "kernel size mismatch");
      ASSERT(kernel_dim + 1 == out_map.coordinate_size(), "kernel size mismatch");
  
      // If either coordinate map is empty
      if (in_map.size() == 0 || out_map.size() == 0) {
        return detail::empty_map_functor<coordinate_type, TemplatedAllocator,
                                         CoordinateMapType, kernel_map_type>()();
      }
  
      if (!is_transpose) {
        if (is_pool && (kernel_stride == kernel_size)) {
          LOG_DEBUG("generating stride_map");
          auto const stride_map =
              detail::stride_map_functor<coordinate_type, TemplatedAllocator,
                                         CoordinateMapType, kernel_map_type>()(
                  in_map, out_map, out_map.get_tensor_stride());
  
          m_kernel_maps[kernel_map_key] = std::move(stride_map);
  
        } else {
          LOG_DEBUG("generating kernel map");
  
          // Default kernel map
          LOG_DEBUG(
              "kernel region with kernel: ",
              PtrToString(kernel_size.data(), in_map.coordinate_size() - 1));
          LOG_DEBUG(
              "kernel region with dilation: ",
              PtrToString(kernel_dilation.data(), in_map.coordinate_size() - 1));
  
          auto kernel_region = cpu_kernel_region<coordinate_type>(
              region_type,                       //
              in_map.coordinate_size(),          //
              in_map.get_tensor_stride().data(), //
              kernel_size.data(),                //
              kernel_dilation.data(),            //
              0, offset.data_ptr<coordinate_type>(), offset.size(0));
  
          auto const kernel_map =
              detail::kernel_map_functor<coordinate_type, TemplatedAllocator,
                                         CoordinateMapType, kernel_map_type>()(
                  in_map, out_map, m_kernel_map_mode, kernel_region);
  
          LOG_DEBUG("kernel_map done");
          m_kernel_maps[kernel_map_key] = std::move(kernel_map);
          LOG_DEBUG("kernel_map saved");
        }
      } else { // is_transpose == true
        // Check first if the out2in kernel map exists
        //
        // Create temporary key for the flipped in/out
        kernel_map_key_type const swapped_kernel_map_key = std::make_tuple(
            p_out_map_key->get_key(), p_in_map_key->get_key(), // maps
            kernel_size, kernel_stride, kernel_dilation,       // kernels
            region_type, false, is_pool);
  
        // Check if the temporary key exists and return swapped in/out
        if (m_kernel_maps.find(swapped_kernel_map_key) != m_kernel_maps.end()) {
          // copy the in out maps from the existing maps
          LOG_DEBUG("found existing kernel_map_key for transposed kernel map");
          m_kernel_maps[kernel_map_key] =
              detail::swap_in_out_map_functor<kernel_map_type>()(
                  m_kernel_maps[swapped_kernel_map_key]);
        } else { // create in out kernel if it doesn't exist
          LOG_DEBUG("No existing kernel_map_key for transposed kernel map");
          if (is_pool && kernel_stride == kernel_size) {
            // e.g. out_map has tensor stride 2 in_map has tensor stride 4.
            // Thus, create a stride map from 2 to 4, out to in.
            auto const stride_map =
                detail::stride_map_functor<coordinate_type, TemplatedAllocator,
                                           CoordinateMapType, kernel_map_type>()(
                    out_map, in_map, in_map.get_tensor_stride());
  
            // TODO Replace the kernel_map values to shared pointers.
            m_kernel_maps[kernel_map_key] =
                detail::swap_in_out_map_functor<kernel_map_type>()(stride_map);
          } else {
            // Default kernel map
            auto kernel_region = cpu_kernel_region<coordinate_type>(
                region_type,                        //
                out_map.coordinate_size(),          //
                out_map.get_tensor_stride().data(), //
                kernel_size.data(),                 //
                kernel_dilation.data(),             //
                0, offset.data_ptr<coordinate_type>(), offset.size(0),
                true // is_transpose
            );
  
            // out to in kernel map
            auto const kernel_map =
                detail::kernel_map_functor<coordinate_type, TemplatedAllocator,
                                           CoordinateMapType, kernel_map_type>()(
                    out_map, in_map, m_kernel_map_mode, kernel_region);
  
            LOG_DEBUG("kernel_map done");
            m_kernel_maps[kernel_map_key] =
                detail::swap_in_out_map_functor<kernel_map_type>()(
                    std::move(kernel_map));
            LOG_DEBUG("kernel_map saved");
          }
        }
      }
    }
  #ifdef DEBUG
    else {
      LOG_DEBUG("kernel map found");
    }
  #endif
  
    // TODO check if it copies or moves the internal data
    return m_kernel_maps[kernel_map_key];
  }
  
  namespace detail {
  
  template <typename coordinate_type>
  struct origin_map_functor<coordinate_type, std::allocator, CoordinateMapCPU,
                            cpu_kernel_map> {
  
    std::pair<at::Tensor, std::vector<at::Tensor>>
    operator()(CoordinateMapCPU<coordinate_type, std::allocator> const
                   &origin_coordinate_map,
               cpu_kernel_map const &origin_map) {
  
      auto options =
          torch::TensorOptions().dtype(torch::kLong).requires_grad(false);
      auto const out_size = origin_coordinate_map.size();
      auto const coordinate_size = origin_coordinate_map.coordinate_size();
  
      at::Tensor batch_indices =
          torch::empty({origin_coordinate_map.size()}, options);
      int64_t *p_batch_indices = batch_indices.data_ptr<int64_t>();
  
      LOG_DEBUG("Copying", origin_coordinate_map.size(), "batch indices");
      for (default_types::index_type i = 0; i < out_size; ++i) {
        p_batch_indices[i] =
            origin_coordinate_map.const_coordinate_data()[i * coordinate_size];
      }
  
      // WARNING: this is an inclusive max index
      coordinate_type const max_batch_index =
          *std::max_element(p_batch_indices, p_batch_indices + out_size);
  
      std::vector<at::Tensor> in_maps;
      for (auto i = 0; i <= max_batch_index; ++i) {
        at::Tensor row_indices = torch::empty({0}, options);
        in_maps.push_back(std::move(row_indices));
      }
  
      ASSERT(origin_map.first.size() == origin_map.second.size(),
             "invalid kernel_map");
      LOG_DEBUG("Iterating over", origin_map.first.size(), "unique maps");
      for (uint32_t out_row_index = 0; out_row_index < origin_map.first.size();
           ++out_row_index) {
        auto const &in_map = origin_map.first[out_row_index];
        int32_t const curr_size = in_map.size();
        ASSERT(curr_size > 0, "invalid kernel map for index", out_row_index);
        auto const curr_batch_index = p_batch_indices[out_row_index];
  
        ASSERT(curr_batch_index <= max_batch_index, "invalid batch index");
        at::Tensor &row_indices = in_maps[curr_batch_index];
        row_indices.resize_({curr_size});
        int64_t *p_row_indices = row_indices.data_ptr<int64_t>();
  
        LOG_DEBUG("Copying", curr_size, "elements to batch index",
                  curr_batch_index, "and row index", out_row_index);
        for (auto i = 0; i < curr_size; ++i) {
          p_row_indices[i] = in_map[i];
        }
      }
  
      return std::make_pair(batch_indices, in_maps);
    }
  };
  
  } // namespace detail
  
  template <typename coordinate_type, typename coordinate_field_type,
            template <typename C> class TemplatedAllocator,
            template <typename T, template <typename Q> class A>
            class CoordinateMapType>
  typename CoordinateMapManager<coordinate_type, coordinate_field_type,
                                TemplatedAllocator,
                                CoordinateMapType>::kernel_map_type const &
  CoordinateMapManager<coordinate_type, coordinate_field_type, TemplatedAllocator,
                       CoordinateMapType>::origin_map(CoordinateMapKey const
                                                          *p_in_map_key) {
    ASSERT(exists(p_in_map_key), ERROR_MAP_NOT_FOUND);
    kernel_map_key_type const kernel_map_key =
        origin_map_key(p_in_map_key->get_key());
    coordinate_map_key_type const origin_key = std::get<1>(kernel_map_key);
  
    if (m_kernel_maps.find(kernel_map_key) == m_kernel_maps.end()) {
      auto const key = origin().first;
      auto const &origin_coordinate_map = m_coordinate_maps.find(key)->second;
      auto origin_map = m_coordinate_maps.find(p_in_map_key->get_key())
                            ->second.origin_map(origin_coordinate_map);
      m_kernel_maps[kernel_map_key] = std::move(origin_map);
    }
  
    return m_kernel_maps[kernel_map_key];
  }
  
  template <typename coordinate_type, typename coordinate_field_type,
            template <typename C> class TemplatedAllocator,
            template <typename T, template <typename Q> class A>
            class CoordinateMapType>
  typename CoordinateMapManager<coordinate_type, coordinate_field_type,
                                TemplatedAllocator,
                                CoordinateMapType>::kernel_map_type const &
  CoordinateMapManager<coordinate_type, coordinate_field_type, TemplatedAllocator,
                       CoordinateMapType>::origin_field_map(CoordinateMapKey const
                                                                *p_in_map_key) {
    ASSERT(exists_field(p_in_map_key), ERROR_MAP_NOT_FOUND);
    kernel_map_key_type const kernel_map_key =
        origin_map_key(p_in_map_key->get_key());
    coordinate_map_key_type const origin_key = std::get<1>(kernel_map_key);
  
    if (m_field_kernel_maps.find(kernel_map_key) == m_field_kernel_maps.end()) {
      auto const key = origin_field().first;
      auto const &origin_coordinate_map = m_coordinate_maps.find(key)->second;
      auto origin_map = m_field_coordinates.find(p_in_map_key->get_key())
                            ->second.origin_map(origin_coordinate_map);
      m_field_kernel_maps[kernel_map_key] = std::move(origin_map);
    }
  
    return m_field_kernel_maps[kernel_map_key];
  }
  namespace detail {
  
  template <typename coordinate_type>
  struct stride_map2tensor_functor<coordinate_type, std::allocator,
                                   CoordinateMapCPU, cpu_kernel_map> {
  
    std::pair<at::Tensor, at::Tensor>
    operator()(cpu_kernel_map const &stride_kernel_map) {
  
      ASSERT(stride_kernel_map.first.size() == 1, "Invalid kernel_map");
      ASSERT(stride_kernel_map.first.size() == stride_kernel_map.second.size(),
             "invalid kernel_map");
  
      auto const &in_map = stride_kernel_map.first[0];
      auto const &out_map = stride_kernel_map.second[0];
  
      auto options =
          torch::TensorOptions().dtype(torch::kLong).requires_grad(false);
      int64_t const out_size = (int64_t)in_map.size();
  
      at::Tensor th_in_map = torch::empty({out_size}, options);
      at::Tensor th_out_map = torch::empty({out_size}, options);
      int64_t *p_in_map = th_in_map.data_ptr<int64_t>();
      int64_t *p_out_map = th_out_map.data_ptr<int64_t>();
  
      // from int32_t to long type
      for (uint32_t i = 0; i < out_size; ++i)
        p_in_map[i] = in_map[i];
      for (uint32_t i = 0; i < out_size; ++i)
        p_out_map[i] = out_map[i];
  
      return std::make_pair(std::move(th_in_map), std::move(th_out_map));
    }
  };
  
  } // namespace detail
  
  template <typename coordinate_type, typename coordinate_field_type,
            template <typename C> class TemplatedAllocator,
            template <typename T, template <typename Q> class A>
            class CoordinateMapType>
  std::pair<at::Tensor, at::Tensor>
  CoordinateMapManager<coordinate_type, coordinate_field_type, TemplatedAllocator,
                       CoordinateMapType>::stride_map_th(CoordinateMapKey const
                                                             *p_in_map_key,
                                                         CoordinateMapKey const
                                                             *p_strided_map_key) {
  
    ASSERT(exists(p_in_map_key), ERROR_MAP_NOT_FOUND);
    ASSERT(exists(p_strided_map_key), ERROR_MAP_NOT_FOUND);
  
    map_type const &in_map =
        m_coordinate_maps.find(p_in_map_key->get_key())->second;
    map_type const &strided_map =
        m_coordinate_maps.find(p_strided_map_key->get_key())->second;
  
    // Get tensor strides and find kernel stride size
    // Check if the kernel map key exists
    auto const &in_map_stride = in_map.get_tensor_stride();
    auto const &strided_map_stride = strided_map.get_tensor_stride();
  
    stride_type kernel_stride(in_map_stride.size());
    for (index_type i = 0; i < kernel_stride.size(); ++i) {
      ASSERT(strided_map_stride[i] % in_map_stride[i] == 0,
             "The tensor stride of the strided map must be divisible by the "
             "tensor stride of the input map. strided_map_stride:",
             ArrToString(strided_map_stride),
             " in_map_stride:", ArrToString(in_map_stride));
      kernel_stride[i] = strided_map_stride[i] / in_map_stride[i];
    }
  
    auto const one_vec = detail::ones(in_map.coordinate_size() - 1);
    kernel_map_key_type const kernel_map_key = std::make_tuple(
        p_in_map_key->get_key(), p_strided_map_key->get_key(), // maps
        kernel_stride, kernel_stride, one_vec,                 // kernels
        RegionType::HYPER_CUBE /* region_type */, 0 /* is_transpose */,
        true /* is_pool */);
  
    if (m_kernel_maps.find(kernel_map_key) == m_kernel_maps.end()) {
      LOG_DEBUG("Creating stride kernel map with kernel size:",
                ArrToString(kernel_stride));
      auto const stride_map =
          detail::stride_map_functor<coordinate_type, TemplatedAllocator,
                                     CoordinateMapType, kernel_map_type>()(
              in_map, strided_map, strided_map.get_tensor_stride());
  
      m_kernel_maps[kernel_map_key] = std::move(stride_map);
    }
  
    // copy the kernel map to tensors
    return detail::stride_map2tensor_functor<coordinate_type, TemplatedAllocator,
                                             CoordinateMapType,
                                             kernel_map_type>()(
        m_kernel_maps[kernel_map_key]);
  }
  
  template <typename coordinate_type, typename coordinate_field_type,
            template <typename C> class TemplatedAllocator,
            template <typename T, template <typename Q> class A>
            class CoordinateMapType>
  std::pair<at::Tensor, std::vector<at::Tensor>>
  CoordinateMapManager<coordinate_type, coordinate_field_type, TemplatedAllocator,
                       CoordinateMapType>::origin_map_th(CoordinateMapKey const
                                                             *p_in_map_key) {
    kernel_map_type const &kernel_map = origin_map(p_in_map_key);
  
    coordinate_map_key_type const origin_key = origin().first;
    map_type const &origin_map = m_coordinate_maps.find(origin_key)->second;
  
    return detail::origin_map_functor<coordinate_type, TemplatedAllocator,
                                      CoordinateMapType, kernel_map_type>()(
        origin_map, kernel_map);
  }
  
  template <typename coordinate_type, typename coordinate_field_type,
            template <typename C> class TemplatedAllocator,
            template <typename T, template <typename Q> class A>
            class CoordinateMapType>
  std::pair<at::Tensor, std::vector<at::Tensor>>
  CoordinateMapManager<coordinate_type, coordinate_field_type, TemplatedAllocator,
                       CoordinateMapType>::
      origin_field_map_th(CoordinateMapKey const *p_in_map_key) {
    kernel_map_type const &kernel_map = origin_field_map(p_in_map_key);
  
    coordinate_map_key_type const origin_key = origin_field().first;
    map_type const &origin_map = m_coordinate_maps.find(origin_key)->second;
  
    return detail::origin_map_functor<coordinate_type, TemplatedAllocator,
                                      CoordinateMapType, kernel_map_type>()(
        origin_map, kernel_map);
  }
  
  // Interpolation map
  template <typename coordinate_type, typename coordinate_field_type,
            template <typename C> class TemplatedAllocator,
            template <typename T, template <typename Q> class A>
            class CoordinateMapType>
  std::vector<at::Tensor>
  CoordinateMapManager<coordinate_type, coordinate_field_type, TemplatedAllocator,
                       CoordinateMapType>::
      interpolation_map_weight(at::Tensor const &tfield,
                               CoordinateMapKey const *p_in_map_key) {
    ASSERT(exists(p_in_map_key), ERROR_MAP_NOT_FOUND);
    return m_coordinate_maps.find(p_in_map_key->get_key())
        ->second.interpolation_map_weight(tfield);
  }
  
  /*********************************/
  /*
  template <typename MapType>
  uint64_t
  CoordsManager<MapType>::createUnionCoords(vector<py::object> py_in_coords_keys,
                                            py::object py_out_coords_key) {
    vector<CoordsKey *> p_in_coords_keys;
    CoordsKey *p_in_coords_key = py_in_coords_keys[0].cast<CoordsKey *>();
    auto tensor_strides = p_in_coords_key->getTensorStride();
    for (const auto &py_in_coords_key : py_in_coords_keys) {
      // Set the tensor strides to the smallest elements.
      p_in_coords_key = py_in_coords_key.cast<CoordsKey *>();
      p_in_coords_keys.push_back(p_in_coords_key);
      transform(tensor_strides.begin(),                            // In1 begin
                tensor_strides.end(),                              // In1 end
                p_in_coords_key->getTensorStride().begin(),        // In2 begin
                tensor_strides.begin(),                            // out begin
                [](int a, int b) -> int { return std::min(a, b); } // binary op
      );
      const uint64_t in_coords_key = p_in_coords_key->getKey();
      ASSERT(existsCoordsKey(in_coords_key),
             "The coord map doesn't exist for the given coords_key: ",
             to_string(in_coords_key), ".");
    }
    CoordsKey *p_out_coords_key = py_out_coords_key.cast<CoordsKey *>();
  
    vector<reference_wrapper<CoordsMap<MapType>>> in_coords_maps;
    for (const CoordsKey *p_in_coords_key : p_in_coords_keys) {
      CoordsMap<MapType> &curr_map = coords_maps[p_in_coords_key->getKey()];
      in_coords_maps.push_back(ref(curr_map));
    }
  
    // set a random coords key
    const uint64_t out_coords_key = getRandomCoordsKey();
  
    // Set the pycoordskey using the last coords_key
    p_out_coords_key->setDimension(p_in_coords_key->getDimension());
    p_out_coords_key->setKey(out_coords_key);
    p_out_coords_key->setTensorStride(tensor_strides);
  
    coords_maps[out_coords_key] =
        CoordsMap<MapType>::union_coords(in_coords_maps);
  
    return out_coords_key;
  }
  
  template <typename MapType>
  const InOutMapKey
  CoordsManager<MapType>::getUnionMapHashKey(vector<py::object> py_in_coords_keys,
                                             py::object py_out_coords_key) const {
    CoordsKey *p_out_coords_key = py_out_coords_key.cast<CoordsKey *>();
    ASSERT(py_in_coords_keys.size() > 1, "Number of input coords must be > 1");
    vector<CoordsKey *> p_in_coords_keys;
    // We use sum of coords key (even with overflow, it will be unique with high
    // prob). We use sum to make the key invariant to the order of the keys.
    uint64_t sum_in_coords_key = 0;
    CoordsKey *p_in_coords_key = py_in_coords_keys[0].cast<CoordsKey *>();
    for (auto &py_in_coords_key : py_in_coords_keys) {
      p_in_coords_key = py_in_coords_key.cast<CoordsKey *>();
      const uint64_t in_coords_key = p_in_coords_key->getKey();
      ASSERT(existsCoordsKey(in_coords_key),
             "The coord map doesn't exist for the given coords_key: ",
             to_string(in_coords_key), ".");
      sum_in_coords_key += in_coords_key;
    }
  
    ASSERT(p_out_coords_key->key_set, "Key is not set. out_coords_key: ",
           to_string(p_out_coords_key->getKey()));
  
    const uint64_t out_coords_key = p_out_coords_key->getKey();
    const vector<int> zero_vec(p_in_coords_key->getDimension(), 0);
    const uint64_t zero_hash = hash_vec(zero_vec);
    InOutMapKey map_key = {sum_in_coords_key,
                           out_coords_key,
                           zero_hash,
                           zero_hash,
                           zero_hash,
                           0,
                           false,
                           true};
    return map_key;
  }
  */
  /**
   * Entry function for coords map generation and the associated kernel maps.
   */
  /*
  template <typename MapType>
  const InOutMapsRefPair<int>
  CoordsManager<MapType>::getPruningInOutMaps(at::Tensor use_feat,
                                              py::object py_in_coords_key,
                                              py::object py_out_coords_key) {
    CoordsKey *p_in_coords_key = py_in_coords_key.cast<CoordsKey *>();
    CoordsKey *p_out_coords_key = py_out_coords_key.cast<CoordsKey *>();
  
    // Create output coordinates if it doesn't exist
    if (!p_out_coords_key->key_set) {
      // The following function setup py_out_coords_key
      createPrunedCoords(use_feat, py_in_coords_key, py_out_coords_key);
    }
  
    const uint64_t in_coords_key = p_in_coords_key->getKey();
    const uint64_t out_coords_key = p_out_coords_key->getKey();
  
    // Use the map key for origin hash map (stride, dilation, kernel are all
    // NULL)
    const InOutMapKey map_key =
        getOriginMapHashKey(py_in_coords_key, py_out_coords_key);
  
    // For non transpose case
    // make a kernel mapping. The kernel will be saved with the map_key.
    if (in_maps.find(map_key) == in_maps.end()) {
      const auto in_out = coords_maps[in_coords_key].pruned_kernel_map(
          coords_maps[out_coords_key]);
      in_maps[map_key] = in_out.first;
      out_maps[map_key] = in_out.second;
    }
  
    return make_pair(ref(in_maps[map_key]), ref(out_maps[map_key]));
  }
  
  template <typename MapType>
  const InOutMapsRefPair<int>
  CoordsManager<MapType>::getUnionInOutMaps(vector<py::object> py_in_coords_keys,
                                            py::object py_out_coords_key) {
    CoordsKey *p_out_coords_key = py_out_coords_key.cast<CoordsKey *>();
  
    // Create output coordinates if it doesn't exist
    if (!p_out_coords_key->key_set)
      createUnionCoords(py_in_coords_keys, py_out_coords_key);
  
    const uint64_t out_coords_key = p_out_coords_key->getKey();
  
    // Map key for origin hash map
    const InOutMapKey map_key =
        getUnionMapHashKey(py_in_coords_keys, py_out_coords_key);
  
    vector<reference_wrapper<CoordsMap<MapType>>> in_coords_maps;
    for (const auto &py_in_coords_key : py_in_coords_keys) {
      const CoordsKey *p_in_coords_key = py_in_coords_key.cast<CoordsKey *>();
      uint64_t in_coords_key = p_in_coords_key->getKey();
      in_coords_maps.push_back(ref(coords_maps[in_coords_key]));
    }
  
    // For non transpose case
    // make a kernel mapping. The kernel will be saved with the map_key.
    if (in_maps.find(map_key) == in_maps.end()) {
      const auto in_out = CoordsMap<MapType>::union_map(
          in_coords_maps, coords_maps[out_coords_key]);
      in_maps[map_key] = in_out.first;
      out_maps[map_key] = in_out.second;
    }
  
    return make_pair(ref(in_maps[map_key]), ref(out_maps[map_key]));
  }
  */
  
  template <typename coordinate_type, typename coordinate_field_type,
            template <typename C> class TemplatedAllocator,
            template <typename T, template <typename Q> class A>
            class CoordinateMapType>
  coordinate_map_key_type
  CoordinateMapManager<coordinate_type, coordinate_field_type, TemplatedAllocator,
                       CoordinateMapType>::
      merge(std::vector<coordinate_map_key_type> const &map_keys) {
    ASSERT(map_keys.size() > 1, "Got one or zero map. Merge at least 2 maps.");
    // Aggregate all coords maps
    std::vector<std::reference_wrapper<map_type>> maps;
    auto const tensor_stride_size = map_keys[0].first.size();
    stride_type merged_map_tensor_stride{map_keys[0].first};
    for (const auto &key : map_keys) {
      ASSERT(exists(key), ERROR_MAP_NOT_FOUND);
      auto &map = m_coordinate_maps.find(key)->second;
      maps.push_back(map);
      for (int k = 0; k < tensor_stride_size; ++k) {
        merged_map_tensor_stride[k] =
            std::min(merged_map_tensor_stride[k], map.get_tensor_stride()[k]);
      }
    }
  
    // Create a merged map with the smallest tensor stride
    coordinate_map_key_type merged_map_key =
        get_random_string_id(merged_map_tensor_stride, "merge");
    map_type const &map = m_coordinate_maps.find(map_keys[0])->second;
    map_type merged_map = map.merge(maps);
    insert(merged_map_key, merged_map);
    return merged_map_key;
  }
  
  template <typename coordinate_type, typename coordinate_field_type,
            template <typename C> class TemplatedAllocator,
            template <typename T, template <typename Q> class A>
            class CoordinateMapType>
  std::pair<coordinate_map_key_type, std::vector<at::Tensor>>
  CoordinateMapManager<coordinate_type, coordinate_field_type, TemplatedAllocator,
                       CoordinateMapType>::
      union_map(std::vector<coordinate_map_key_type> const &map_keys) {
    // Create a merged map
    auto const merged_key = merge(map_keys);
    map_type const &merged_map = m_coordinate_maps.find(merged_key)->second;
  
    std::vector<std::reference_wrapper<map_type>> maps;
    for (const auto &key : map_keys) {
      ASSERT(exists(key), ERROR_MAP_NOT_FOUND);
      maps.push_back(std::ref(m_coordinate_maps.find(key)->second));
    }
  
    return std::make_pair(merged_key, merged_map.union_map(maps));
  }
  
  template <typename coordinate_type, typename coordinate_field_type,
            template <typename C> class TemplatedAllocator,
            template <typename T, template <typename Q> class A>
            class CoordinateMapType>
  std::vector<at::Tensor>
  CoordinateMapManager<coordinate_type, coordinate_field_type, TemplatedAllocator,
                       CoordinateMapType>::
      union_map_th(std::vector<CoordinateMapKey *> const &p_map_keys,
                   CoordinateMapKey *p_out_key) {
    ASSERT(!p_out_key->is_key_set(),
           "Out coordinate map key should be uninitialized");
  
    std::vector<coordinate_map_key_type> map_keys;
    map_keys.reserve(p_map_keys.size());
    std::for_each(
        p_map_keys.begin(), p_map_keys.end(),
        [&](CoordinateMapKey *p_key) { map_keys.push_back(p_key->get_key()); });
    auto union_pair = union_map(map_keys);
    p_out_key->set_key(union_pair.first);
    return union_pair.second;
  }
  
  /* Helper functions */
  template <typename coordinate_type, typename coordinate_field_type,
            template <typename C> class TemplatedAllocator,
            template <typename T, template <typename Q> class A>
            class CoordinateMapType>
  at::Tensor
  CoordinateMapManager<coordinate_type, coordinate_field_type, TemplatedAllocator,
                       CoordinateMapType>::get_coordinates(CoordinateMapKey const
                                                               *p_key) const {
    ASSERT(exists(p_key), ERROR_MAP_NOT_FOUND);
    auto const it = m_coordinate_maps.find(p_key->get_key());
    ASSERT(it != m_coordinate_maps.end(), ERROR_MAP_NOT_FOUND);
    auto const &map = it->second;
    auto const nrows = map.size();
    auto const ncols = map.coordinate_size();
    LOG_DEBUG("coordinate map nrows:", nrows, "ncols:", ncols);
  
    // CPU torch.IntTensor
    auto options = torch::TensorOptions().dtype(torch::kInt).requires_grad(false);
    if (!detail::is_cpu_coordinate_map<CoordinateMapType>::value) {
  #ifndef CPU_ONLY
      auto device_id = at::cuda::current_device();
      options = options.device(torch::kCUDA, device_id);
  #else
      ASSERT(false, ERROR_CPU_ONLY);
  #endif
    }
    at::Tensor coordinates = torch::empty({(long)nrows, (long)ncols}, options);
  
    LOG_DEBUG("Initialized coordinates");
    // copy to the out coords
    map.copy_coordinates(coordinates.template data_ptr<coordinate_type>());
    LOG_DEBUG("Copied coordinates");
    return coordinates;
  }
  
  namespace detail {
  
  template <typename coordinate_type>
  struct kernel_map_to_tensors<coordinate_type, std::allocator, CoordinateMapCPU,
                               cpu_kernel_map> {
  
    std::unordered_map<int64_t, at::Tensor>
    operator()(cpu_kernel_map const &kernel_map) {
  
      const auto &in_maps = kernel_map.first;
      const auto &out_maps = kernel_map.second;
  
      auto options =
          torch::TensorOptions().dtype(torch::kInt).requires_grad(false);
  
      std::unordered_map<int64_t, at::Tensor> th_kernel_maps;
      for (auto k = 0; k < in_maps.size(); ++k) {
        const auto &in_map = in_maps[k];
        const auto &out_map = out_maps[k];
        const long N = in_map.size();
        if (N > 0) {
          at::Tensor kernel_map = torch::empty({2, N}, options);
          int32_t *p_kernel_map = kernel_map.data_ptr<int32_t>();
          std::copy_n(&in_map[0], N, p_kernel_map);
          std::copy_n(&out_map[0], N, p_kernel_map + N);
          th_kernel_maps[k] = std::move(kernel_map);
        }
      }
  
      return th_kernel_maps;
    }
  };
  
  } // namespace detail
  
  template <typename coordinate_type, typename coordinate_field_type,
            template <typename C> class TemplatedAllocator,
            template <typename T, template <typename Q> class A>
            class CoordinateMapType>
  std::unordered_map<int64_t, at::Tensor> CoordinateMapManager<
      coordinate_type, coordinate_field_type, TemplatedAllocator,
      CoordinateMapType>::kernel_map_th(CoordinateMapKey const *p_in_map_key,
                                        CoordinateMapKey const *p_out_map_key,
                                        stride_type const &kernel_size, //
                                        stride_type const &kernel_stride,
                                        stride_type const &kernel_dilation,
                                        RegionType::Type const region_type,
                                        at::Tensor const &offset,
                                        bool is_transpose, bool is_pool) {
  
    auto const &curr_kernel_map =
        kernel_map(p_in_map_key, p_out_map_key,                 // maps
                   kernel_size, kernel_stride, kernel_dilation, // kernels
                   region_type, offset, is_transpose, is_pool);
  
    return detail::kernel_map_to_tensors<coordinate_type, TemplatedAllocator,
                                         CoordinateMapType, kernel_map_type>()(
        curr_kernel_map);
  }
  
  template <typename coordinate_type, typename coordinate_field_type,
            template <typename C> class TemplatedAllocator,
            template <typename T, template <typename Q> class A>
            class CoordinateMapType>
  at::Tensor CoordinateMapManager<coordinate_type, coordinate_field_type,
                                  TemplatedAllocator, CoordinateMapType>::
      get_coordinate_field(CoordinateMapKey const *p_key) const {
    auto const it = m_field_coordinates.find(p_key->get_key());
    ASSERT(it != m_field_coordinates.end(), ERROR_MAP_NOT_FOUND);
    auto const &map = it->second;
    auto const nrows = map.size();
    auto const ncols = map.coordinate_size();
  
    auto options = torch::TensorOptions()
                       .dtype(std::is_same<float, coordinate_field_type>::value
                                  ? torch::kFloat
                                  : torch::kDouble)
                       .requires_grad(false);
  
    if (!detail::is_cpu_coordinate_map<CoordinateMapType>::value) {
  #ifndef CPU_ONLY
      auto device_id = at::cuda::current_device();
      options = options.device(torch::kCUDA, device_id);
  #else
      ASSERT(false, ERROR_CPU_ONLY);
  #endif
    }
    at::Tensor coordinates = torch::empty({(long)nrows, (long)ncols}, options);
  
    // copy to the out coords
    map.copy_coordinates(coordinates.template data_ptr<coordinate_field_type>());
    return coordinates;
  }
  
  template class CoordinateMapManager<default_types::dcoordinate_type,
                                      default_types::ccoordinate_type,
                                      std::allocator, CoordinateMapCPU>;
  
  } // end namespace minkowski
  