numpy/core/numeric.pyi

import sys
from typing import (
    Any,
    Optional,
    Union,
    Sequence,
    Tuple,
    Callable,
    List,
    overload,
    TypeVar,
    Iterable,
)

from numpy import ndarray, generic, dtype, bool_, signedinteger, _OrderKACF, _OrderCF
from numpy.typing import ArrayLike, DTypeLike, _ShapeLike

if sys.version_info >= (3, 8):
    from typing import Literal
else:
    from typing_extensions import Literal

_T = TypeVar("_T")
_ArrayType = TypeVar("_ArrayType", bound=ndarray)

_CorrelateMode = Literal["valid", "same", "full"]

@overload
def zeros_like(
    a: _ArrayType,
    dtype: None = ...,
    order: _OrderKACF = ...,
    subok: Literal[True] = ...,
    shape: None = ...,
) -> _ArrayType: ...
@overload
def zeros_like(
    a: ArrayLike,
    dtype: DTypeLike = ...,
    order: _OrderKACF = ...,
    subok: bool = ...,
    shape: Optional[_ShapeLike] = ...,
) -> ndarray: ...
def ones(
    shape: _ShapeLike,
    dtype: DTypeLike = ...,
    order: _OrderCF = ...,
    *,
    like: ArrayLike = ...,
) -> ndarray: ...
@overload
def ones_like(
    a: _ArrayType,
    dtype: None = ...,
    order: _OrderKACF = ...,
    subok: Literal[True] = ...,
    shape: None = ...,
) -> _ArrayType: ...
@overload
def ones_like(
    a: ArrayLike,
    dtype: DTypeLike = ...,
    order: _OrderKACF = ...,
    subok: bool = ...,
    shape: Optional[_ShapeLike] = ...,
) -> ndarray: ...
@overload
def empty_like(
    a: _ArrayType,
    dtype: None = ...,
    order: _OrderKACF = ...,
    subok: Literal[True] = ...,
    shape: None = ...,
) -> _ArrayType: ...
@overload
def empty_like(
    a: ArrayLike,
    dtype: DTypeLike = ...,
    order: _OrderKACF = ...,
    subok: bool = ...,
    shape: Optional[_ShapeLike] = ...,
) -> ndarray: ...
def full(
    shape: _ShapeLike,
    fill_value: Any,
    dtype: DTypeLike = ...,
    order: _OrderCF = ...,
    *,
    like: ArrayLike = ...,
) -> ndarray: ...
@overload
def full_like(
    a: _ArrayType,
    fill_value: Any,
    dtype: None = ...,
    order: _OrderKACF = ...,
    subok: Literal[True] = ...,
    shape: None = ...,
) -> _ArrayType: ...
@overload
def full_like(
    a: ArrayLike,
    fill_value: Any,
    dtype: DTypeLike = ...,
    order: _OrderKACF = ...,
    subok: bool = ...,
    shape: Optional[_ShapeLike] = ...,
) -> ndarray: ...
@overload
def count_nonzero(
    a: ArrayLike, axis: None = ..., *, keepdims: Literal[False] = ...
) -> int: ...
@overload
def count_nonzero(
    a: ArrayLike, axis: _ShapeLike = ..., *, keepdims: bool = ...
) -> Union[signedinteger[Any], ndarray]: ...  # TODO: np.intp
def isfortran(a: Union[ndarray, generic]) -> bool: ...
def argwhere(a: ArrayLike) -> ndarray: ...
def flatnonzero(a: ArrayLike) -> ndarray: ...
def correlate(a: ArrayLike, v: ArrayLike, mode: _CorrelateMode = ...) -> ndarray: ...
def convolve(a: ArrayLike, v: ArrayLike, mode: _CorrelateMode = ...) -> ndarray: ...
@overload
def outer(a: ArrayLike, b: ArrayLike, out: None = ...) -> ndarray: ...
@overload
def outer(a: ArrayLike, b: ArrayLike, out: _ArrayType = ...) -> _ArrayType: ...
def tensordot(
    a: ArrayLike,
    b: ArrayLike,
    axes: Union[int, Tuple[_ShapeLike, _ShapeLike]] = ...,
) -> ndarray: ...
def roll(
    a: ArrayLike,
    shift: _ShapeLike,
    axis: Optional[_ShapeLike] = ...,
) -> ndarray: ...
def rollaxis(a: ndarray, axis: int, start: int = ...) -> ndarray: ...
def moveaxis(
    a: ndarray,
    source: _ShapeLike,
    destination: _ShapeLike,
) -> ndarray: ...
def cross(
    a: ArrayLike,
    b: ArrayLike,
    axisa: int = ...,
    axisb: int = ...,
    axisc: int = ...,
    axis: Optional[int] = ...,
) -> ndarray: ...
@overload
def indices(
    dimensions: Sequence[int],
    dtype: DTypeLike = ...,
    sparse: Literal[False] = ...,
) -> ndarray: ...
@overload
def indices(
    dimensions: Sequence[int],
    dtype: DTypeLike = ...,
    sparse: Literal[True] = ...,
) -> Tuple[ndarray, ...]: ...
def fromfunction(
    function: Callable[..., _T],
    shape: Sequence[int],
    *,
    dtype: DTypeLike = ...,
    like: ArrayLike = ...,
    **kwargs: Any,
) -> _T: ...
def isscalar(element: Any) -> bool: ...
def binary_repr(num: int, width: Optional[int] = ...) -> str: ...
def base_repr(number: int, base: int = ..., padding: int = ...) -> str: ...
def identity(n: int, dtype: DTypeLike = ..., *, like: ArrayLike = ...) -> ndarray: ...
def allclose(
    a: ArrayLike,
    b: ArrayLike,
    rtol: float = ...,
    atol: float = ...,
    equal_nan: bool = ...,
) -> bool: ...
def isclose(
    a: ArrayLike,
    b: ArrayLike,
    rtol: float = ...,
    atol: float = ...,
    equal_nan: bool = ...,
) -> Union[bool_, ndarray]: ...
def array_equal(a1: ArrayLike, a2: ArrayLike) -> bool: ...
def array_equiv(a1: ArrayLike, a2: ArrayLike) -> bool: ...