from __future__ import annotations

from typing import Any, Optional

def Print(
    input: Any,
    first_n: int = ...,
    message: Optional[Any] = ...,
    summarize: int = ...,
    print_tensor_name: bool = ...,
    print_tensor_type: bool = ...,
    print_tensor_shape: bool = ...,
    print_tensor_layout: bool = ...,
    print_tensor_lod: bool = ...,
    print_phase: str = ...,
): ...
def Assert(cond: Any, data: Optional[Any] = ..., summarize: int = ..., name: Optional[Any] = ...): ...

class BlockGuard:
    main_program: Any = ...
    def __init__(self, main_program: Any) -> None: ...
    def __enter__(self) -> None: ...
    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any): ...

class BlockGuardWithCompletion(BlockGuard):
    rnn: Any = ...
    def __init__(self, rnn: Any) -> None: ...
    def __enter__(self): ...
    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any): ...

class StaticRNNMemoryLink:
    init: Any = ...
    pre_mem: Any = ...
    mem: Any = ...
    def __init__(self, init: Any, pre_mem: Any, mem: Optional[Any] = ...) -> None: ...

class StaticRNN:
    BEFORE_RNN_BLOCK: int = ...
    IN_RNN_BLOCK: int = ...
    AFTER_RNN_BLOCK: int = ...
    helper: Any = ...
    memories: Any = ...
    inputs: Any = ...
    outputs: Any = ...
    status: Any = ...
    seq_len: Any = ...
    def __init__(self, name: Optional[Any] = ...) -> None: ...
    def step(self): ...
    def memory(
        self,
        init: Optional[Any] = ...,
        shape: Optional[Any] = ...,
        batch_ref: Optional[Any] = ...,
        init_value: float = ...,
        init_batch_dim_idx: int = ...,
        ref_batch_dim_idx: int = ...,
    ): ...
    def step_input(self, x: Any): ...
    def step_output(self, o: Any) -> None: ...
    def output(self, *outputs: Any) -> None: ...
    def update_memory(self, mem: Any, var: Any) -> None: ...
    def __call__(self, *args: Any, **kwargs: Any): ...

class WhileGuard(BlockGuard):
    while_op: Any = ...
    def __init__(self, while_op: Any) -> None: ...
    def __enter__(self): ...
    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any): ...

class While:
    BEFORE_WHILE_BLOCK: int = ...
    IN_WHILE_BLOCK: int = ...
    AFTER_WHILE_BLOCK: int = ...
    helper: Any = ...
    status: Any = ...
    cond_var: Any = ...
    is_test: Any = ...
    def __init__(self, cond: Any, is_test: bool = ..., name: Optional[Any] = ...): ...
    def block(self): ...

def while_loop(cond: Any, body: Any, loop_vars: Any, is_test: bool = ..., name: Optional[Any] = ...): ...
def increment(x: Any, value: float = ..., in_place: bool = ...): ...
def array_write(x: Any, i: Any, array: Optional[Any] = ...): ...
def create_array(dtype: Any, initialized_list: Optional[Any] = ...): ...
def less_than(x: Any, y: Any, force_cpu: Optional[Any] = ..., cond: Optional[Any] = ..., name: Optional[Any] = ...): ...
def less_equal(x: Any, y: Any, cond: Optional[Any] = ..., name: Optional[Any] = ...): ...
def greater_than(x: Any, y: Any, cond: Optional[Any] = ..., name: Optional[Any] = ...): ...
def greater_equal(x: Any, y: Any, cond: Optional[Any] = ..., name: Optional[Any] = ...): ...
def equal(x: Any, y: Any, cond: Optional[Any] = ..., name: Optional[Any] = ...): ...
def not_equal(x: Any, y: Any, cond: Optional[Any] = ..., name: Optional[Any] = ...): ...
def array_read(array: Any, i: Any): ...
def array_length(array: Any): ...

class ConditionalBlockGuard(BlockGuard):
    block: Any = ...
    def __init__(self, block: Any) -> None: ...
    def __enter__(self): ...
    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any): ...

class ConditionalBlock:
    inputs: Any = ...
    is_scalar_condition: Any = ...
    helper: Any = ...
    def __init__(self, inputs: Any, is_scalar_condition: bool = ..., name: Optional[Any] = ...) -> None: ...
    def block(self): ...
    def complete(self) -> None: ...
    def need_append_conditional_block_grad(self, inside_block: Any): ...
    def append_conditional_block_grad(
        self, parent_block: Any, inside_block: Any, conditional_block_op: Any
    ) -> None: ...

def cond(pred: Any, true_fn: Optional[Any] = ..., false_fn: Optional[Any] = ..., name: Optional[Any] = ...): ...
def case(pred_fn_pairs: Any, default: Optional[Any] = ..., name: Optional[Any] = ...): ...

class Switch:
    helper: Any = ...
    inside_scope: bool = ...
    pre_not_conditions: Any = ...
    def __init__(self, name: Optional[Any] = ...) -> None: ...
    def case(self, condition: Any): ...
    def default(self): ...
    def __enter__(self): ...
    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any): ...

class IfElseBlockGuard:
    is_true: Any = ...
    ie: Any = ...
    cond_block: Any = ...
    def __init__(self, is_true: Any, ifelse: Any) -> None: ...
    def __enter__(self) -> None: ...
    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any): ...

class IfElse:
    OUT_IF_ELSE_BLOCKS: int = ...
    IN_IF_ELSE_TRUE_BLOCKS: int = ...
    IN_IF_ELSE_FALSE_BLOCKS: int = ...
    helper: Any = ...
    cond: Any = ...
    input_table: Any = ...
    status: Any = ...
    conditional_true_block: Any = ...
    conditional_false_block: Any = ...
    output_table: Any = ...
    def __init__(self, cond: Any, name: Optional[Any] = ...) -> None: ...
    def input(self, x: Any): ...
    def true_block(self): ...
    def false_block(self): ...
    def output(self, *outs: Any) -> None: ...
    def __call__(self): ...

class DynamicRNN:
    BEFORE_RNN: int = ...
    IN_RNN: int = ...
    AFTER_RNN: int = ...
    helper: Any = ...
    status: Any = ...
    lod_rank_table: Any = ...
    max_seq_len: Any = ...
    step_idx: Any = ...
    zero_idx: Any = ...
    mem_dict: Any = ...
    output_array: Any = ...
    outputs: Any = ...
    cond: Any = ...
    while_op: Any = ...
    input_array: Any = ...
    mem_link: Any = ...
    def __init__(self, name: Optional[Any] = ...) -> None: ...
    def step_input(self, x: Any, level: int = ...): ...
    def static_input(self, x: Any): ...
    def block(self) -> None: ...
    def __call__(self, *args: Any, **kwargs: Any): ...
    def memory(
        self,
        init: Optional[Any] = ...,
        shape: Optional[Any] = ...,
        value: float = ...,
        need_reorder: bool = ...,
        dtype: str = ...,
    ): ...
    def update_memory(self, ex_mem: Any, new_mem: Any) -> None: ...
    def output(self, *outputs: Any) -> None: ...

def switch_case(branch_index: Any, branch_fns: Any, default: Optional[Any] = ..., name: Optional[Any] = ...): ...
def reorder_lod_tensor_by_rank(x: Any, rank_table: Any): ...
def is_empty(x: Any, name: Optional[Any] = ...): ...
