from __future__ import annotations
from typing import Iterable, Optional, TypeVar
from collections import OrderedDict
from nova.nn.parameter import (
    UninitializedParameter,
    UninitializedBuffer,
    Parameter,
    Buffer,
)
from nova import Tensor

TModule = TypeVar("TModule", bound="Module")

def _addindent(s_: str, numSpaces: int) -> str: ...

class ModuleMeta(type):
    def __new__(mcs, name, bases, namespace, **kwargs): ...

class Module(metaclass=ModuleMeta):
    _initialized: bool
    _parameters: dict[str, Parameter]
    _buffers: dict[str, Buffer]
    _modules: dict[str, Module]
    _training: bool

    def __init__(self) -> None: ...
    def __call__(self, *args, **kwargs) -> Tensor: ...
    def forward(self, *args, **kwargs) -> Tensor: ...
    def parameters(self, recurse: bool = True) -> Iterable[Parameter]: ...
    def buffers(self, recurse: bool = True) -> Iterable[Buffer]: ...
    def modules(self, recurse: bool = True) -> Iterable[Module]: ...
    def register_buffer(self, name: str, buffer: Optional[Buffer]) -> None: ...
    def register_parameter(self, name: str, param: Optional[Parameter]) -> None: ...
    def register_module(self, name: str, module: Optional[Module]) -> None: ...
    def __setattr__(
        self,
        name: str,
        value: (
            Parameter | Module | Buffer | UninitializedBuffer | UninitializedParameter
        ),
    ) -> None: ...
    def named_parameters(
        self, prefix: str = "", recurse: bool = True
    ) -> Iterable[tuple[str, Parameter]]: ...
    def named_buffers(
        self, prefix: str = "", recurse: bool = True
    ) -> Iterable[tuple[str, Buffer]]: ...
    def named_modules(self, prefix: str = "") -> Iterable[tuple[str, Module]]: ...
    def train(self: TModule, mode: bool = True) -> TModule: ...
    def eval(self: TModule, mode: bool = False) -> TModule: ...
    def state_dict(
        self, destination: Optional[OrderedDict | dict] = None, prefix: str = ""
    ) -> dict: ...
    def load_state_dict(
        self, state_dict: OrderedDict | dict, prefix: str = ""
    ) -> None: ...
    def __repr__(self) -> str: ...
    def _get_name(self) -> str: ...
    def extra_repr(self) -> str: ...
