from __future__ import annotations

from typing import Any, Optional

from paddle.fluid.dygraph import layers

def prepare_context(strategy: Optional[Any] = ...): ...

class ParallelEnv:
    def __init__(self) -> None: ...
    @property
    def rank(self): ...
    @property
    def world_size(self): ...
    @property
    def device_id(self): ...
    @property
    def current_endpoint(self): ...
    @property
    def trainer_endpoints(self): ...
    @property
    def nrings(self): ...
    local_rank: Any = ...
    nranks: Any = ...
    dev_id: Any = ...

Env = ParallelEnv

class DataParallel(layers.Layer):
    find_unused_parameters: Any = ...
    grad_need_sync: bool = ...
    group: Any = ...
    var_dtype: Any = ...
    comm_buffer_size: Any = ...
    last_comm_buffer_size: Any = ...
    def __init__(
        self,
        layers: Any,
        strategy: Optional[Any] = ...,
        comm_buffer_size: int = ...,
        last_comm_buffer_size: int = ...,
        find_unused_parameters: bool = ...,
        group: Optional[Any] = ...,
    ) -> None: ...
    group_indices: Any = ...
    def init_reducer(self): ...
    def no_sync(self) -> None: ...
    def forward(self, *inputs: Any, **kwargs: Any): ...
    def scale_loss(self, loss: Any): ...
    def apply_collective_grads(self) -> None: ...
    def state_dict(
        self, destination: Optional[Any] = ..., include_sublayers: bool = ..., structured_name_prefix: str = ...
    ): ...
    def set_state_dict(self, state_dict: Any, use_structured_name: bool = ...) -> None: ...
    set_dict: Any = ...
    load_dict: Any = ...
