from __future__ import annotations

from typing import Any

from paddle.distributed.fleet.meta_optimizers.sharding.utils import *

class Shard:
    global_params: Any = ...
    worker_idx: int = ...
    worker_num: int = ...
    global_param2device: Any = ...
    device2global_params: Any = ...
    def __init__(self) -> None: ...
    def setup(self, params_grads: Any, worker_idx: Any, worker_num: Any) -> None: ...
    def has_param(self, var_name: Any): ...
    def has_opt_var(self, var_name: Any): ...
    def has_var(self, var_name: Any): ...
    def find_broadcast_params(self, block: Any): ...
    def device(self, var_name: Any): ...
    def is_param(self, var_name: Any): ...
    def is_opti_var(self, var_name: Any): ...
    def filter_grads(self, grads: Any): ...

class ProgramSegment:
    def __init__(self, block: Any) -> None: ...
