from __future__ import annotations

from collections import OrderedDict as OrderedDict
from typing import Any, Optional

from .pass_base import PassBase as PassBase
from .pass_base import PassType as PassType
from .pass_base import register_pass as register_pass

def find_adjacent_match_sequences(iterable: Any, filter_func: Any, adjacent_filter_func: Optional[Any] = ...): ...
def insert_fuse_all_reduce_ops(
    block: Any, reversed_op_indices: Any, input_var_names: Any, output_var_names: Any, dtype: Any, attrs: Any
): ...
def has_same_attrs(op1: Any, op2: Any, attr_names: Any): ...
def filter_all_collective_op_indices(block: Any): ...
def find_all_fuse_all_reduce_groups(block: Any): ...
def split_fuse_all_reduce_groups_by_deps(block: Any, groups: Any, op_deps: Any): ...
def insert_coalesce_tensor_ops(block: Any, coalesce_ops_kwargs: Any): ...
def insert_fuse_all_reduce_by_memory_size(block: Any, groups: Any, max_memory_size: Any) -> None: ...

class FuseAllReducePass(PassBase):
    def __init__(self) -> None: ...
