def _nndescent_knn_graph(
    x,
    k,
    segs,
    num_iters=None,
    max_candidates=None,
    delta=0.001,
    sample_rate=0.5,
    dist="euclidean",
):
    r"""Construct multiple graphs from multiple sets of points according to
    **approximate** k-nearest-neighbor using NN-descent algorithm from paper
    `Efficient k-nearest neighbor graph construction for generic similarity
    measures <https://www.cs.princeton.edu/cass/papers/www11.pdf>`_.

    Parameters
    ----------
    x : Tensor
        Coordinates/features of points. Must be 2D. It can be either on CPU or GPU.
    k : int
        The number of nearest neighbors per node.
    segs : list[int]
        Number of points in each point set. The numbers in :attr:`segs`
        must sum up to the number of rows in :attr:`x`.
    num_iters : int, optional
        The maximum number of NN-descent iterations to perform. A value will be
        chosen based on the size of input by default.
        (Default: None)
    max_candidates : int, optional
        The maximum number of candidates to be considered during one iteration.
        Larger values will provide more accurate search results later, but
        potentially at non-negligible computation cost. A value will be chosen
        based on the number of neighbors by default.
        (Default: None)
    delta : float, optional
        A value controls the early abort. This function will abort if
        :math:`k * N * delta > c`, where :math:`N` is the number of points,
        :math:`c` is the number of updates during last iteration.
        (Default: 0.001)
    sample_rate : float, optional
        A value controls how many candidates sampled. It should be a float value
        between 0 and 1. Larger values will provide higher accuracy and converge
        speed but with higher time cost.
        (Default: 0.5)
    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
    -------
    DGLGraph
        The graph. The node IDs are in the same order as :attr:`x`.
    """
    num_points, _ = F.shape(x)
    if isinstance(segs, (tuple, list)):
        segs = F.tensor(segs)
    segs = F.copy_to(segs, F.context(x))

    if max_candidates is None:
        max_candidates = min(60, k)
    if num_iters is None:
        num_iters = max(10, int(round(np.log2(num_points))))
    max_candidates = int(sample_rate * max_candidates)

    # 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)

    # k must less than or equal to min(segs)
    if k > F.min(segs, dim=0):
        raise DGLError(
            "'k' must be less than or equal to the number of points in 'x'"
            "expect 'k' <= {}, got 'k' = {}".format(F.min(segs, dim=0), k)
        )
    if delta < 0 or delta > 1:
        raise DGLError("'delta' must in [0, 1], got 'delta' = {}".format(delta))

    offset = F.zeros((F.shape(segs)[0] + 1,), F.dtype(segs), F.context(segs))
    offset[1:] = F.cumsum(segs, dim=0)
    out = F.zeros((2, num_points * k), F.dtype(segs), F.context(segs))

    # points, offsets, out, k, num_iters, max_candidates, delta
    _CAPI_DGLNNDescent(
        F.to_dgl_nd(x),
        F.to_dgl_nd(offset),
        F.zerocopy_to_dgl_ndarray_for_write(out),
        k,
        num_iters,
        max_candidates,
        delta,
    )
    return out
Construct multiple graphs from multiple sets of points according to
    **approximate** k-nearest-neighbor using NN-descent algorithm from paper
    `Efficient k-nearest neighbor graph construction for generic similarity
    measures <https://www.cs.princeton.edu/cass/papers/www11.pdf>`_.

    Parameters
    ----------
    x : Tensor
        Coordinates/features of points. Must be 2D. It can be either on CPU or GPU.
    k : int
        The number of nearest neighbors per node.
    segs : list[int]
        Number of points in each point set. The numbers in :attr:`segs`
        must sum up to the number of rows in :attr:`x`.
    num_iters : int, optional
        The maximum number of NN-descent iterations to perform. A value will be
        chosen based on the size of input by default.
        (Default: None)
    max_candidates : int, optional
        The maximum number of candidates to be considered during one iteration.
        Larger values will provide more accurate search results later, but
        potentially at non-negligible computation cost. A value will be chosen
        based on the number of neighbors by default.
        (Default: None)
    delta : float, optional
        A value controls the early abort. This function will abort if
        :math:`k * N * delta > c`, where :math:`N` is the number of points,
        :math:`c` is the number of updates during last iteration.
        (Default: 0.001)
    sample_rate : float, optional
        A value controls how many candidates sampled. It should be a float value
        between 0 and 1. Larger values will provide higher accuracy and converge
        speed but with higher time cost.
        (Default: 0.5)
    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
    -------
    DGLGraph
        The graph. The node IDs are in the same order as :attr:`x`.
    
python/dgl/transforms/functional.py