from typing import Any, Literal, Optional, Union
from nova.nn import Parameter, Buffer
from nova import Tensor
from nova._typing import Size

def calculate_gain(
    nonlinearity: Literal[
        "linear", "sigmoid", "tanh", "relu", "leaky_relu", "prelu", "gelu"
    ],
    param: Optional[float] = None,
) -> float: ...
def _validate_mode(mode: str) -> None: ...
def _calculate_fans(shape: Size) -> tuple[int, int]: ...
def get_fans(
    tensor: Tensor, mode: Literal["both", "fan_in", "fan_out"] = "fan_in"
) -> Union[int, tuple[int, int]]: ...
def xavier_normal_(tensor: Parameter | Buffer, gain: float = 1.0) -> None: ...
def xavier_uniform_(tensor: Parameter | Buffer, gain: float = 1.0) -> None: ...
def kaiming_normal_(
    tensor: Parameter | Buffer,
    a: Optional[float] = None,
    nonlinearity: str = "leaky_relu",
    mode: str = "fan_in",
) -> None: ...
def kaiming_uniform_(
    tensor: Parameter | Buffer,
    a: Optional[float] = None,
    nonlinearity: str = "relu",
    mode: str = "fan_in",
) -> None: ...
def uniform_(
    tensor: Tensor | Parameter | Buffer, low: float = 0, high: float = 1
) -> None: ...
def normal_(
    tensor: Tensor | Parameter | Buffer, mean: float = 0, std: float = 1
) -> None: ...
def constant_(tensor: Tensor | Parameter | Buffer, val: Any) -> None: ...
def zeros_(tensor: Tensor | Parameter | Buffer) -> None: ...
def ones_(tensor: Tensor | Parameter | Buffer) -> None: ...
def random_(tensor: Tensor | Parameter | Buffer) -> None: ...
