from __future__ import annotations

from typing import Any

from paddle.distributed.auto_parallel.process_mesh import ProcessMesh as ProcessMesh
from paddle.fluid.backward import ProgramStats

from .pass_base import PassBase as PassBase
from .pass_base import register_pass as register_pass

class RecomputeState(ProgramStats):
    var_op_deps: Any = ...
    def __init__(self, block: Any, ops: Any) -> None: ...
    def build_stats(self) -> None: ...
    def get_recompute_segments(self, checkpoints: Any): ...
    def modify_forward_desc_for_recompute(self, dist_context: Any) -> None: ...

class RecomputePass(PassBase):
    def __init__(self) -> None: ...
    def reset_op_dist_attr(self, op: Any, var_name_dict: Any) -> None: ...
    def set_op_dist_attr(self, op: Any, old_dist_attr: Any, var_name_dict: Any) -> None: ...
