from typing_extensions import TypeIs
from abc import abstractmethod
from typing import TYPE_CHECKING, Optional
from nova import Tensor

if TYPE_CHECKING:
    from nova._typing import Dtype, Size

class Parameter(Tensor):
    is_bn_param: bool
    def __init__(
        self,
        data: Tensor = ...,
        dtype: Optional[Dtype] = None,
        requires_grad: bool = ...,
    ) -> None: ...

class Buffer(Tensor):
    persistent: bool
    def __init__(
        self,
        data: Tensor = ...,
        dtype: Optional[Dtype] = None,
        *,
        persistent: bool = ...,
    ) -> None: ...

class UninitializedParameter(UninitializedTensorMixin, Parameter):
    def __init__(self, requires_grad: bool = ...) -> None: ...
    def materialize(self, shape: Size, dtype: Optional[Dtype] = None) -> Parameter: ...

class UninitializedBuffer(UninitializedTensorMixin, Buffer):
    def __init__(self, persistent: bool = ...) -> None: ...
    def materialize(self, shape: Size, dtype: Optional[Dtype] = None) -> Buffer: ...

class UninitializedTensorMixin:
    @abstractmethod
    def materialize(self, shape: Size, dtype: Optional[Dtype] = None) -> Tensor: ...
    @property
    def shape(self): ...
    def __repr__(self) -> str: ...

def is_lazy(
    param: Tensor,
) -> TypeIs[UninitializedParameter | UninitializedBuffer]: ...
