"""
This type stub file was generated by pyright.
"""

"""Module for degree bucketing schedulers."""
def gen_degree_bucketing_schedule(reduce_udf, message_ids, dst_nodes, recv_nodes, var_nf, var_mf, var_out, ntype=...): # -> None:
    """Create degree bucketing schedule.

    The messages will be divided by their receivers into buckets. Each bucket
    contains nodes that have the same in-degree. The reduce UDF will be applied
    on each bucket. The per-bucket result will be merged according to the
    *unique-ascending order* of the recv node ids. The order is important to
    be compatible with other reduce scheduler such as v2v_spmv.

    Parameters
    ----------
    reduce_udf : callable
        The UDF to reduce messages.
    message_ids : utils.Index
        The variable for message ids.
        Invariant: len(message_ids) == len(dst_nodes)
    dst_nodes : utils.Index
        The variable for dst node of each message.
        Invariant: len(message_ids) == len(dst_nodes)
    recv_nodes : utils.Index
        The unique nodes that perform recv.
        Invariant: recv_nodes = sort(unique(dst_nodes))
    var_nf : var.FEAT_DICT
        The variable for node feature frame.
    var_mf : var.FEAT_DICT
        The variable for message frame.
    var_out : var.FEAT_DICT
        The variable for output feature dicts.
    ntype : str, optional
        The node type, if running on a heterograph.
        If None, assuming it's running on a homogeneous graph.
    """
    ...

def gen_group_apply_edge_schedule(apply_func, u, v, eid, group_by, var_src_nf, var_dst_nf, var_ef, var_out, canonical_etype=...): # -> None:
    """Create degree bucketing schedule for group_apply_edge

    Edges will be grouped by either its source node or destination node
    specified by 'group_by', and will be divided into buckets in which
    'group_by' nodes have the same degree. The apply_func UDF will be applied
    to each bucket. The per-bucket result will be merged according to the
    *unique-ascending order* of the edge ids.

    Parameters
    ----------
    apply_func: callable
        The edge_apply_func UDF
    u: utils.Index
        Source nodes of edges to apply
    v: utils.Index
        Destination nodes of edges to apply
    eid: utils.Index
        Edges to apply
    group_by: str
        If "src", group by u. If "dst", group by v
    var_src_nf : var.FEAT_DICT
        The variable for source feature frame.
    var_dst_nf : var.FEAT_DICT
        The variable for destination feature frame.
    var_ef : var.FEAT_DICT
        The variable for edge frame.
    var_out : var.FEAT_DICT
        The variable for output feature dicts.
    canonical_etype : tuple[str, str, str], optional
        Canonical edge type if running on a heterograph.
        Default: (None, None, None), if running on a homogeneous graph.
    """
    ...

