from __future__ import annotations

from typing import Any, Optional

class _DecoderType:
    TRAINING: int = ...
    BEAM_SEARCH: int = ...

class InitState:
    def __init__(
        self,
        init: Optional[Any] = ...,
        shape: Optional[Any] = ...,
        value: float = ...,
        init_boot: Optional[Any] = ...,
        need_reorder: bool = ...,
        dtype: str = ...,
    ) -> None: ...
    @property
    def value(self): ...
    @property
    def need_reorder(self): ...

class _MemoryState:
    def __init__(self, state_name: Any, rnn_obj: Any, init_state: Any) -> None: ...
    def get_state(self): ...
    def update_state(self, state: Any) -> None: ...

class _ArrayState:
    def __init__(self, state_name: Any, block: Any, init_state: Any) -> None: ...
    def get_state(self): ...
    def update_state(self, state: Any) -> None: ...

class StateCell:
    def __init__(self, inputs: Any, states: Any, out_state: Any, name: Optional[Any] = ...) -> None: ...
    def get_state(self, state_name: Any): ...
    def get_input(self, input_name: Any): ...
    def set_state(self, state_name: Any, state_value: Any) -> None: ...
    def state_updater(self, updater: Any): ...
    def compute_state(self, inputs: Any) -> None: ...
    def update_states(self) -> None: ...
    def out_state(self): ...

class TrainingDecoder:
    BEFORE_DECODER: int = ...
    IN_DECODER: int = ...
    AFTER_DECODER: int = ...
    def __init__(self, state_cell: Any, name: Optional[Any] = ...) -> None: ...
    def block(self) -> None: ...
    @property
    def state_cell(self): ...
    @property
    def dynamic_rnn(self): ...
    @property
    def type(self): ...
    def step_input(self, x: Any): ...
    def static_input(self, x: Any): ...
    def __call__(self, *args: Any, **kwargs: Any): ...
    def output(self, *outputs: Any) -> None: ...

class BeamSearchDecoder:
    BEFORE_BEAM_SEARCH_DECODER: int = ...
    IN_BEAM_SEARCH_DECODER: int = ...
    AFTER_BEAM_SEARCH_DECODER: int = ...
    def __init__(
        self,
        state_cell: Any,
        init_ids: Any,
        init_scores: Any,
        target_dict_dim: Any,
        word_dim: Any,
        input_var_dict: Any = ...,
        topk_size: int = ...,
        sparse_emb: bool = ...,
        max_len: int = ...,
        beam_size: int = ...,
        end_id: int = ...,
        name: Optional[Any] = ...,
    ) -> None: ...
    def block(self) -> None: ...
    @property
    def type(self): ...
    def early_stop(self) -> None: ...
    def decode(self) -> None: ...
    def read_array(self, init: Any, is_ids: bool = ..., is_scores: bool = ...): ...
    def update_array(self, array: Any, value: Any) -> None: ...
    def __call__(self): ...
    @property
    def state_cell(self): ...
