from __future__ import annotations

from typing import Any

from .auto_parallel_amp import AMPPass as AMPPass
from .pass_base import register_pass as register_pass

world_process_group: Any
__amp_skip_ops__: Any

def set_op_dtype_to_fp16(op: Any) -> None: ...

class FP16State:
    program: Any = ...
    amp_list: Any = ...
    use_fp16_guard: Any = ...
    dist_context: Any = ...
    grad_op_to_op_map: Any = ...
    forward_non_leaf_tensors: Any = ...
    forward_input_cast_ops: Any = ...
    is_train: bool = ...
    def __init__(self, program: Any, amp_list: Any, dist_context: Any, use_fp16_guard: Any) -> None: ...
    def set_var_to_fp16(self, var_name: Any, block: Any) -> None: ...
    def resolute_tensor_dtype(self, block: Any) -> None: ...
    def cast_block(self, block: Any) -> None: ...

class FP16Pass(AMPPass):
    def __init__(self) -> None: ...
