from __future__ import annotations

from collections import OrderedDict as OrderedDict
from typing import Any

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

OpRole: Any
OP_ROLE_KEY: Any

class ShardingPass(PassBase):
    dp_groups: Any = ...
    sharding_infos: Any = ...
    varname_to_sharding_info: Any = ...
    partial_sharding: bool = ...
    outer_dp_group: Any = ...
    def __init__(self) -> None: ...

def shard_parameters(params: Any, group_size: Any): ...

class ShardingInfo:
    group: Any = ...
    params: Any = ...
    param_names: Any = ...
    group_size: Any = ...
    global_rank: Any = ...
    local_rank: Any = ...
    rank_to_params: Any = ...
    param_to_rank: Any = ...
    def __init__(self, group: Any, rank: Any, params: Any) -> None: ...
    def get_var_rank(self, varname: Any): ...
    def is_in_local_shard(self, param_name: Any): ...
    def get_broadcast_vars_and_param_usage(self, block: Any): ...
