from __future__ import annotations

from typing import Any

from paddle.distributed.fleet.meta_optimizers.sharding.utils import *

class FP16Utils:
    def __init__(self) -> None: ...
    @staticmethod
    def is_fp16_cast_op(block: Any, op: Any, params: Any): ...
    @staticmethod
    def is_fp32_cast_op(block: Any, op: Any): ...
    @staticmethod
    def remove_cast_op(block: Any, params: Any, segment: Any, offset: Any): ...
    @staticmethod
    def prune_fp16(block: Any, shard: Any, reduced_grads_to_param: Any, ring_ids: Any) -> None: ...
    @staticmethod
    def sync_amp_check_nan_inf(block: Any, ring_ids: Any) -> None: ...
