code: def knn(
      k,
      x,
      x_segs,
      y=None,
      y_segs=None,
      algorithm="bruteforce",
      dist="euclidean",
  ):
      r"""For each element in each segment in :attr:`y`, find :attr:`k` nearest
      points in the same segment in :attr:`x`. If :attr:`y` is None, perform a self-query
      over :attr:`x`.
  
      This function allows multiple point sets with different capacity. The points
      from different sets are stored contiguously in the :attr:`x` and :attr:`y` tensor.
      :attr:`x_segs` and :attr:`y_segs` specifies the number of points in each point set.
  
      Parameters
      ----------
      k : int
          The number of nearest neighbors per node.
      x : Tensor
          The point coordinates in x. It can be either on CPU or GPU (must be the
          same as :attr:`y`). Must be 2D.
      x_segs : Union[List[int], Tensor]
          Number of points in each point set in :attr:`x`. The numbers in :attr:`x_segs`
          must sum up to the number of rows in :attr:`x`.
      y : Tensor, optional
          The point coordinates in y. It can be either on CPU or GPU (must be the
          same as :attr:`x`). Must be 2D.
          (default: None)
      y_segs : Union[List[int], Tensor], optional
          Number of points in each point set in :attr:`y`. The numbers in :attr:`y_segs`
          must sum up to the number of rows in :attr:`y`.
          (default: None)
      algorithm : str, optional
          Algorithm used to compute the k-nearest neighbors.
  
          * 'bruteforce' will compute distances pair by pair and
            directly select the k-nearest neighbors during distance
            computation. This method is slower than 'bruteforce-blas'
            but has less memory overhead (i.e., :math:`O(Nk)` where :math:`N`
            is the number of points, :math:`k` is the number of nearest
            neighbors per node) since we do not need to store all distances.
  
          * 'bruteforce-sharemem' (CUDA only) is similar to 'bruteforce'
            but use shared memory in CUDA devices for buffer. This method is
            faster than 'bruteforce' when the dimension of input points
            is not large. This method is only available on CUDA device.
  
          * 'kd-tree' will use the kd-tree algorithm (CPU only).
            This method is suitable for low-dimensional data (e.g. 3D
            point clouds)
  
          * 'nn-descent' is an approximate approach from paper
            `Efficient k-nearest neighbor graph construction for generic similarity
            measures <https://www.cs.princeton.edu/cass/papers/www11.pdf>`_. This method
            will search for nearest neighbor candidates in "neighbors' neighbors".
  
          Note: Currently, 'nn-descent' only supports self-query cases, i.e. :attr:`y` is None.
          (default: 'bruteforce')
      dist : str, optional
          The distance metric used to compute distance between points. It can be the following
          metrics:
          * 'euclidean': Use Euclidean distance (L2 norm) :math:`\sqrt{\sum_{i} (x_{i} - y_{i})^{2}}`.
          * 'cosine': Use cosine distance.
          (default: 'euclidean')
  
      Returns
      -------
      Tensor
          Tensor with size `(2, k * num_points(y))`
          The first subtensor contains point indexs in :attr:`y`. The second subtensor contains
          point indexs in :attr:`x`
      """
      # TODO(lygztq) add support for querying different point sets using nn-descent.
      if algorithm == "nn-descent":
          if y is not None or y_segs is not None:
              raise DGLError(
                  "Currently 'nn-descent' only supports self-query cases."
              )
          return _nndescent_knn_graph(x, k, x_segs, dist=dist)
  
      # self query
      if y is None:
          y = x
          y_segs = x_segs
  
      assert F.context(x) == F.context(y)
      if isinstance(x_segs, (tuple, list)):
          x_segs = F.tensor(x_segs)
      if isinstance(y_segs, (tuple, list)):
          y_segs = F.tensor(y_segs)
      x_segs = F.copy_to(x_segs, F.context(x))
      y_segs = F.copy_to(y_segs, F.context(y))
  
      # k shoule be less than or equal to min(x_segs)
      min_num_points = F.min(x_segs, dim=0)
      if k > min_num_points:
          dgl_warning(
              "'k' should be less than or equal to the number of points in 'x'"
              "expect k <= {0}, got k = {1}, use k = {0}".format(
                  min_num_points, k
              )
          )
          k = F.as_scalar(min_num_points)
  
      # invalid k
      if k <= 0:
          raise DGLError(
              "Invalid k value. expect k > 0, got k = {}".format(k)
          )
  
      # empty point set
      if F.shape(x)[0] == 0 or F.shape(y)[0] == 0:
          raise DGLError("Find empty point set")
  
      dist = dist.lower()
      dist_metric_list = ["euclidean", "cosine"]
      if dist not in dist_metric_list:
          raise DGLError(
              "Only {} are supported for distance"
              "computation, got {}".format(dist_metric_list, dist)
          )
  
      x_offset = F.zeros(
          (F.shape(x_segs)[0] + 1,), F.dtype(x_segs), F.context(x_segs)
      )
      x_offset[1:] = F.cumsum(x_segs, dim=0)
      y_offset = F.zeros(
          (F.shape(y_segs)[0] + 1,), F.dtype(y_segs), F.context(y_segs)
      )
      y_offset[1:] = F.cumsum(y_segs, dim=0)
  
      out = F.zeros(
          (2, F.shape(y)[0] * k), F.dtype(x_segs), F.context(x_segs)
      )
  
      # if use cosine distance, normalize input points first
      # thus we can use euclidean distance to find knn equivalently.
      if dist == "cosine":
          l2_norm = lambda v: F.sqrt(F.sum(v * v, dim=1, keepdims=True))
          x = x / (l2_norm(x) + 1e-5)
          y = y / (l2_norm(y) + 1e-5)
  
      _CAPI_DGLKNN(
          F.to_dgl_nd(x),
          F.to_dgl_nd(x_offset),
          F.to_dgl_nd(y),
          F.to_dgl_nd(y_offset),
          k,
          F.zerocopy_to_dgl_ndarray_for_write(out),
          algorithm,
      )
      return out
  
imports: import copy
  from collections.abc import Iterable, Mapping
  import numpy
  import scipy.sparse
  import scipy.sparse.linalg
  from ..utils import version
  from ..None import backend, batch, convert, function, ndarray, subgraph, utils
  from .._ffi.function import _init_api
  from ..base import dgl_warning, DGLError, EID, NID
  from ..frame import Frame
  from ..heterograph import DGLGraph
  from ..heterograph_index import create_heterograph_from_relations, create_metagraph_index
  from ..partition import metis_partition, metis_partition_assignment, partition_graph_with_halo
  from ..sampling.neighbor import sample_neighbors
label_desc: Use image node_id 7 for calling a global function with example usage: knn(k, x, x_segs, y, y_segs, algorithm, dist) and returns: out, _nndescent_knn_graph