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
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`
    
python/dgl/transforms/functional.py