from __future__ import annotations

from typing import Any

class Collective:
    nrings: Any = ...
    endpoints: Any = ...
    current_endpoint: Any = ...
    other_endpoints: Any = ...
    nranks: Any = ...
    rank: Any = ...
    startup_program: Any = ...
    main_program: Any = ...
    op_role_key: Any = ...
    op_role_var_key: Any = ...
    def __init__(self, nrings: Any) -> None: ...
    wait_port: Any = ...
    def transpile(
        self, startup_program: Any, main_program: Any, rank: Any, endpoints: Any, current_endpoint: Any, wait_port: Any
    ) -> None: ...

class GradAllReduce(Collective):
    mode: str = ...
    def __init__(self, nrings: int = ...) -> None: ...

class LocalSGD(Collective):
    snapshot_key: str = ...
    mode: str = ...
    def __init__(self, nrings: int = ...) -> None: ...
    def snapshot_name(self, param_name: Any): ...

class SingleProcessMultiThread(GradAllReduce):
    mode: str = ...
    def __init__(self) -> None: ...

class MultiThread(GradAllReduce):
    mode: str = ...
    trans_mode: Any = ...
    fuse_grad_size_in_num: int = ...
    gpu_num: Any = ...
    def __init__(self, nrings: int = ..., trans_mode: str = ...) -> None: ...
