from __future__ import annotations

from typing import Any, Optional

from paddle.fluid.dygraph.layers import Layer

from ...utils.log_util import layer_to_str as layer_to_str
from ...utils.log_util import logger as logger

class LayerDesc:
    layer_func: Any = ...
    inputs: Any = ...
    kwargs: Any = ...
    def __init__(self, layer_func: Any, *inputs: Any, **kwargs: Any) -> None: ...
    def build_layer(self): ...

class SharedLayerDesc(LayerDesc):
    layer_name: Any = ...
    forward_func: Any = ...
    shared_weight_attr: Any = ...
    def __init__(
        self,
        key: Any,
        layer_func: Any,
        forward_func: Optional[Any] = ...,
        shared_weight_attr: str = ...,
        *inputs: Any,
        **kwargs: Any,
    ) -> None: ...

class SegmentLayers:
    method: Any = ...
    num_parts: Any = ...
    num_items: Any = ...
    def __init__(self, layers_desc: Any, num_parts: Any, method: str = ...) -> None: ...
    def do_segment(self): ...
    def uniform(self, num_items: Any, num_parts: Any): ...

class PipelineLayer(Layer):
    device_id: Any = ...
    layers: Any = ...
    global_rank: Any = ...
    shared_layers: Any = ...
    shared_weight_attrs: Any = ...
    run_function: Any = ...
    shared_comm: Any = ...
    def __init__(
        self,
        layers: Any,
        num_stages: Optional[Any] = ...,
        topology: Optional[Any] = ...,
        loss_fn: Optional[Any] = ...,
        seg_method: str = ...,
        recompute_interval: int = ...,
        recompute_offload: bool = ...,
        recompute_partition: bool = ...,
    ) -> None: ...
    def get_stage_from_index(self, layer_idx: Any): ...
    def allreduce_shared_weight_gradients(self) -> None: ...
    def forward_function(self, start: Any, end: Any): ...
    def forward(self, input: Any): ...
    def save_state_dict(self, path: Any): ...
    def set_state_dir(self, path: Any) -> None: ...
