from __future__ import annotations

from collections import abc as abc
from enum import Enum
from math import inf as inf
from typing import Any

class Taskflow:
    task: Any = ...
    callback: Any = ...
    def __init__(self, task: Any, callback: Any) -> None: ...

class Type(Enum):
    fp16: Any = ...
    fp32: Any = ...

class ShardingClipGrad:
    def __init__(self, clip: Any, device: Any, group: Any) -> None: ...
    def __getattr__(self, item: Any): ...
    def __call__(self, params_grads: Any): ...

def device_guard(dev_id: int = ..., device: str = ...) -> None: ...
def ShardingScaler(scaler: Any): ...
