from __future__ import annotations

from typing import Any, Optional

from ..meta_parallel.sharding.sharding_utils import Type as Type
from ..meta_parallel.sharding.sharding_utils import device_guard as device_guard

class InternalStorage:
    buffer: Any = ...
    def __init__(self, size: Any, dtype: Any, device: Any, convert_cpu: bool = ...) -> None: ...
    def to(self, device: Any, dtype: Optional[Any] = ..., keep_alignment: bool = ...) -> None: ...

class ParamStorage(InternalStorage):
    param2align: Any = ...
    def __init__(self, size: Any, dtype: Any, device: Any) -> None: ...
    def to(self, device: Any, dtype: Optional[Any] = ..., keep_alignment: bool = ...) -> None: ...
    buffer: Any = ...
    def add_rank_params(self, trainable_params: Any, param2align: Any, convert_gpu: bool = ...) -> None: ...

class GradStorage(InternalStorage):
    params_checked_in: int = ...
    destination: Any = ...
    sent: bool = ...
    def __init__(
        self, size: Any, dtype: Any, device: Any, destination: Any, parm2align: Any, convert_cpu: bool = ...
    ) -> None: ...
    def reset_checked_in(self) -> None: ...
    @property
    def all_checked_in(self): ...
    def can_add_grad_view(self, param: Any, align: Any): ...
    def to(self, device: Any, dtype: Optional[Any] = ..., keep_alignment: bool = ...) -> None: ...
    def add_grad(self, param: Any, align: Any) -> None: ...
    buffer: Any = ...
    def manumal_relase(self) -> None: ...
    def rebuild(self) -> None: ...
