from __future__ import annotations

from functools import partial as partial
from typing import Any, Optional

from paddle.device import get_cudnn_version as get_cudnn_version
from paddle.device import get_device as get_device
from paddle.fluid.data_feeder import convert_dtype as convert_dtype
from paddle.fluid.layers.utils import pack_sequence_as as pack_sequence_as
from paddle.nn import Layer, LayerList

def split_states(states: Any, bidirectional: bool = ..., state_components: int = ...): ...
def concat_states(states: Any, bidirectional: bool = ..., state_components: int = ...): ...

class RNNCellBase(Layer):
    shape: Any = ...
    def get_initial_states(
        self,
        batch_ref: Any,
        shape: Optional[Any] = ...,
        dtype: Optional[Any] = ...,
        init_value: float = ...,
        batch_dim_idx: int = ...,
    ): ...
    @property
    def state_shape(self) -> None: ...
    @property
    def state_dtype(self) -> None: ...

class SimpleRNNCell(RNNCellBase):
    weight_ih: Any = ...
    weight_hh: Any = ...
    bias_ih: Any = ...
    bias_hh: Any = ...
    input_size: Any = ...
    hidden_size: Any = ...
    activation: Any = ...
    def __init__(
        self,
        input_size: Any,
        hidden_size: Any,
        activation: str = ...,
        weight_ih_attr: Optional[Any] = ...,
        weight_hh_attr: Optional[Any] = ...,
        bias_ih_attr: Optional[Any] = ...,
        bias_hh_attr: Optional[Any] = ...,
        name: Optional[Any] = ...,
    ) -> None: ...
    def forward(self, inputs: Any, states: Optional[Any] = ...): ...
    @property
    def state_shape(self): ...
    def extra_repr(self): ...

class LSTMCell(RNNCellBase):
    weight_ih: Any = ...
    weight_hh: Any = ...
    bias_ih: Any = ...
    bias_hh: Any = ...
    hidden_size: Any = ...
    input_size: Any = ...
    def __init__(
        self,
        input_size: Any,
        hidden_size: Any,
        weight_ih_attr: Optional[Any] = ...,
        weight_hh_attr: Optional[Any] = ...,
        bias_ih_attr: Optional[Any] = ...,
        bias_hh_attr: Optional[Any] = ...,
        name: Optional[Any] = ...,
    ) -> None: ...
    def forward(self, inputs: Any, states: Optional[Any] = ...): ...
    @property
    def state_shape(self): ...
    def extra_repr(self): ...

class GRUCell(RNNCellBase):
    weight_ih: Any = ...
    weight_hh: Any = ...
    bias_ih: Any = ...
    bias_hh: Any = ...
    hidden_size: Any = ...
    input_size: Any = ...
    def __init__(
        self,
        input_size: Any,
        hidden_size: Any,
        weight_ih_attr: Optional[Any] = ...,
        weight_hh_attr: Optional[Any] = ...,
        bias_ih_attr: Optional[Any] = ...,
        bias_hh_attr: Optional[Any] = ...,
        name: Optional[Any] = ...,
    ) -> None: ...
    def forward(self, inputs: Any, states: Optional[Any] = ...): ...
    @property
    def state_shape(self): ...
    def extra_repr(self): ...

class RNN(Layer):
    cell: Any = ...
    is_reverse: Any = ...
    time_major: Any = ...
    def __init__(self, cell: Any, is_reverse: bool = ..., time_major: bool = ...) -> None: ...
    def forward(
        self, inputs: Any, initial_states: Optional[Any] = ..., sequence_length: Optional[Any] = ..., **kwargs: Any
    ): ...

class BiRNN(Layer):
    cell_fw: Any = ...
    cell_bw: Any = ...
    time_major: Any = ...
    def __init__(self, cell_fw: Any, cell_bw: Any, time_major: bool = ...) -> None: ...
    def forward(
        self, inputs: Any, initial_states: Optional[Any] = ..., sequence_length: Optional[Any] = ..., **kwargs: Any
    ): ...

class RNNBase(LayerList):
    mode: Any = ...
    input_size: Any = ...
    hidden_size: Any = ...
    dropout: Any = ...
    num_directions: Any = ...
    time_major: Any = ...
    num_layers: Any = ...
    state_components: Any = ...
    could_use_cudnn: bool = ...
    def __init__(
        self,
        mode: Any,
        input_size: Any,
        hidden_size: Any,
        num_layers: int = ...,
        direction: str = ...,
        time_major: bool = ...,
        dropout: float = ...,
        weight_ih_attr: Optional[Any] = ...,
        weight_hh_attr: Optional[Any] = ...,
        bias_ih_attr: Optional[Any] = ...,
        bias_hh_attr: Optional[Any] = ...,
    ) -> None: ...
    def flatten_parameters(self) -> None: ...
    def forward(self, inputs: Any, initial_states: Optional[Any] = ..., sequence_length: Optional[Any] = ...): ...
    def extra_repr(self): ...

class SimpleRNN(RNNBase):
    activation: Any = ...
    def __init__(
        self,
        input_size: Any,
        hidden_size: Any,
        num_layers: int = ...,
        direction: str = ...,
        time_major: bool = ...,
        dropout: float = ...,
        activation: str = ...,
        weight_ih_attr: Optional[Any] = ...,
        weight_hh_attr: Optional[Any] = ...,
        bias_ih_attr: Optional[Any] = ...,
        bias_hh_attr: Optional[Any] = ...,
        name: Optional[Any] = ...,
    ) -> None: ...

class LSTM(RNNBase):
    def __init__(
        self,
        input_size: Any,
        hidden_size: Any,
        num_layers: int = ...,
        direction: str = ...,
        time_major: bool = ...,
        dropout: float = ...,
        weight_ih_attr: Optional[Any] = ...,
        weight_hh_attr: Optional[Any] = ...,
        bias_ih_attr: Optional[Any] = ...,
        bias_hh_attr: Optional[Any] = ...,
        name: Optional[Any] = ...,
    ) -> None: ...

class GRU(RNNBase):
    def __init__(
        self,
        input_size: Any,
        hidden_size: Any,
        num_layers: int = ...,
        direction: str = ...,
        time_major: bool = ...,
        dropout: float = ...,
        weight_ih_attr: Optional[Any] = ...,
        weight_hh_attr: Optional[Any] = ...,
        bias_ih_attr: Optional[Any] = ...,
        bias_hh_attr: Optional[Any] = ...,
        name: Optional[Any] = ...,
    ) -> None: ...
