from typing import Literal

from jax import Array

from . import linalg

def mean(arr: Array, axis: int = -1) -> Array: ...
def fill_diagonal(
    arr: Array, val: Array | float, *, inplace: Literal[False]
) -> Array: ...
def exp(arr: Array) -> Array: ...
def abs(arr: Array) -> Array: ...  # noqa: A001
def triu_indices(n: int, k: int = 0) -> tuple[Array, Array]: ...
def argpartition(arr: Array, ind: int) -> Array: ...
def log(arr: Array) -> Array: ...
def array(obj: object) -> Array: ...
def arccosh(obj: object) -> Array: ...
def allclose(
    left: Array, right: Array, rtol: float = ..., atol: float = ...
) -> Array: ...

inf: float

__all__ = ("linalg",)
