from __future__ import annotations

from typing import Any, Optional

from .utils import *

class RNNCell:
    def call(self, inputs: Any, states: Any, **kwargs: Any) -> None: ...
    def __call__(self, inputs: Any, states: Any, **kwargs: Any): ...
    shape: Any = ...
    def get_initial_states(
        self,
        batch_ref: Any,
        shape: Optional[Any] = ...,
        dtype: str = ...,
        init_value: int = ...,
        batch_dim_idx: int = ...,
    ): ...
    @property
    def state_shape(self) -> None: ...
    @property
    def state_dtype(self) -> None: ...

class GRUCell(RNNCell):
    hidden_size: Any = ...
    gru_unit: Any = ...
    def __init__(
        self,
        hidden_size: Any,
        param_attr: Optional[Any] = ...,
        bias_attr: Optional[Any] = ...,
        gate_activation: Optional[Any] = ...,
        activation: Optional[Any] = ...,
        dtype: str = ...,
        name: str = ...,
    ) -> None: ...
    def call(self, inputs: Any, states: Any): ...
    @property
    def state_shape(self): ...

class LSTMCell(RNNCell):
    hidden_size: Any = ...
    lstm_unit: Any = ...
    def __init__(
        self,
        hidden_size: Any,
        param_attr: Optional[Any] = ...,
        bias_attr: Optional[Any] = ...,
        gate_activation: Optional[Any] = ...,
        activation: Optional[Any] = ...,
        forget_bias: float = ...,
        dtype: str = ...,
        name: str = ...,
    ) -> None: ...
    def call(self, inputs: Any, states: Any): ...
    @property
    def state_shape(self): ...

def rnn(
    cell: Any,
    inputs: Any,
    initial_states: Optional[Any] = ...,
    sequence_length: Optional[Any] = ...,
    time_major: bool = ...,
    is_reverse: bool = ...,
    **kwargs: Any,
): ...

class ArrayWrapper:
    array: Any = ...
    def __init__(self, x: Any) -> None: ...
    def append(self, x: Any): ...
    def __getitem__(self, item: Any): ...

def birnn(
    cell_fw: Any,
    cell_bw: Any,
    inputs: Any,
    initial_states: Optional[Any] = ...,
    sequence_length: Optional[Any] = ...,
    time_major: bool = ...,
    **kwargs: Any,
): ...

class Decoder:
    def initialize(self, inits: Any) -> None: ...
    def step(self, time: Any, inputs: Any, states: Any, **kwargs: Any) -> None: ...
    def finalize(self, outputs: Any, final_states: Any, sequence_lengths: Any) -> None: ...
    @property
    def tracks_own_finished(self): ...

class BeamSearchDecoder(Decoder):
    cell: Any = ...
    embedding_fn: Any = ...
    output_fn: Any = ...
    start_token: Any = ...
    end_token: Any = ...
    beam_size: Any = ...
    def __init__(
        self,
        cell: Any,
        start_token: Any,
        end_token: Any,
        beam_size: Any,
        embedding_fn: Optional[Any] = ...,
        output_fn: Optional[Any] = ...,
    ) -> None: ...
    @staticmethod
    def tile_beam_merge_with_batch(x: Any, beam_size: Any): ...

    class OutputWrapper: ...
    class StateWrapper: ...
    kinf: float = ...
    batch_size: Any = ...
    start_token_tensor: Any = ...
    end_token_tensor: Any = ...
    def initialize(self, initial_cell_states: Any): ...
    def step(self, time: Any, inputs: Any, states: Any, **kwargs: Any): ...
    def finalize(self, outputs: Any, final_states: Any, sequence_lengths: Any): ...
    @property
    def tracks_own_finished(self): ...

def dynamic_decode(
    decoder: Any,
    inits: Optional[Any] = ...,
    max_step_num: Optional[Any] = ...,
    output_time_major: bool = ...,
    impute_finished: bool = ...,
    is_test: bool = ...,
    return_length: bool = ...,
    **kwargs: Any,
): ...

class DecodeHelper:
    def initialize(self) -> None: ...
    def sample(self, time: Any, outputs: Any, states: Any) -> None: ...
    def next_inputs(self, time: Any, outputs: Any, states: Any, sample_ids: Any) -> None: ...

class TrainingHelper(DecodeHelper):
    inputs: Any = ...
    sequence_length: Any = ...
    time_major: Any = ...
    inputs_: Any = ...
    def __init__(self, inputs: Any, sequence_length: Any, time_major: bool = ...): ...
    def initialize(self): ...
    def sample(self, time: Any, outputs: Any, states: Any): ...
    def next_inputs(self, time: Any, outputs: Any, states: Any, sample_ids: Any): ...

class GreedyEmbeddingHelper(DecodeHelper):
    embedding_fn: Any = ...
    start_tokens: Any = ...
    end_token: Any = ...
    def __init__(self, embedding_fn: Any, start_tokens: Any, end_token: Any) -> None: ...
    def initialize(self): ...
    def sample(self, time: Any, outputs: Any, states: Any): ...
    def next_inputs(self, time: Any, outputs: Any, states: Any, sample_ids: Any): ...

class SampleEmbeddingHelper(GreedyEmbeddingHelper):
    softmax_temperature: Any = ...
    seed: Any = ...
    def __init__(
        self,
        embedding_fn: Any,
        start_tokens: Any,
        end_token: Any,
        softmax_temperature: Optional[Any] = ...,
        seed: Optional[Any] = ...,
    ) -> None: ...
    def sample(self, time: Any, outputs: Any, states: Any): ...

class BasicDecoder(Decoder):
    cell: Any = ...
    helper: Any = ...
    output_fn: Any = ...
    def __init__(self, cell: Any, helper: Any, output_fn: Optional[Any] = ...) -> None: ...
    def initialize(self, initial_cell_states: Any): ...

    class OutputWrapper: ...

    def step(self, time: Any, inputs: Any, states: Any, **kwargs: Any): ...

def dynamic_lstm(
    input: Any,
    size: Any,
    h_0: Optional[Any] = ...,
    c_0: Optional[Any] = ...,
    param_attr: Optional[Any] = ...,
    bias_attr: Optional[Any] = ...,
    use_peepholes: bool = ...,
    is_reverse: bool = ...,
    gate_activation: str = ...,
    cell_activation: str = ...,
    candidate_activation: str = ...,
    dtype: str = ...,
    name: Optional[Any] = ...,
): ...
def lstm(
    input: Any,
    init_h: Any,
    init_c: Any,
    max_len: Any,
    hidden_size: Any,
    num_layers: Any,
    dropout_prob: float = ...,
    is_bidirec: bool = ...,
    is_test: bool = ...,
    name: Optional[Any] = ...,
    default_initializer: Optional[Any] = ...,
    seed: int = ...,
): ...
def dynamic_lstmp(
    input: Any,
    size: Any,
    proj_size: Any,
    param_attr: Optional[Any] = ...,
    bias_attr: Optional[Any] = ...,
    use_peepholes: bool = ...,
    is_reverse: bool = ...,
    gate_activation: str = ...,
    cell_activation: str = ...,
    candidate_activation: str = ...,
    proj_activation: str = ...,
    dtype: str = ...,
    name: Optional[Any] = ...,
    h_0: Optional[Any] = ...,
    c_0: Optional[Any] = ...,
    cell_clip: Optional[Any] = ...,
    proj_clip: Optional[Any] = ...,
): ...
def dynamic_gru(
    input: Any,
    size: Any,
    param_attr: Optional[Any] = ...,
    bias_attr: Optional[Any] = ...,
    is_reverse: bool = ...,
    gate_activation: str = ...,
    candidate_activation: str = ...,
    h_0: Optional[Any] = ...,
    origin_mode: bool = ...,
): ...
def gru_unit(
    input: Any,
    hidden: Any,
    size: Any,
    param_attr: Optional[Any] = ...,
    bias_attr: Optional[Any] = ...,
    activation: str = ...,
    gate_activation: str = ...,
    origin_mode: bool = ...,
): ...
def beam_search(
    pre_ids: Any,
    pre_scores: Any,
    ids: Any,
    scores: Any,
    beam_size: Any,
    end_id: Any,
    level: int = ...,
    is_accumulated: bool = ...,
    name: Optional[Any] = ...,
    return_parent_idx: bool = ...,
): ...
def beam_search_decode(ids: Any, scores: Any, beam_size: Any, end_id: Any, name: Optional[Any] = ...): ...
def lstm_unit(
    x_t: Any,
    hidden_t_prev: Any,
    cell_t_prev: Any,
    forget_bias: float = ...,
    param_attr: Optional[Any] = ...,
    bias_attr: Optional[Any] = ...,
    name: Optional[Any] = ...,
): ...
