import builtins
from collections.abc import Callable, Sequence
from typing import (
Any,
TypeAlias,
overload,
TypeVar,
Literal as L,
)
import numpy as np
from numpy import (
generic,
number,
timedelta64,
datetime64,
int_,
intp,
float64,
complex128,
signedinteger,
floating,
complexfloating,
object_,
_OrderCF,
)
from numpy._typing import (
DTypeLike,
_DTypeLike,
ArrayLike,
_ArrayLike,
NDArray,
_SupportsArray,
_SupportsArrayFunc,
_ArrayLikeInt_co,
_ArrayLikeFloat_co,
_ArrayLikeComplex_co,
_ArrayLikeObject_co,
)
_T = TypeVar("_T")
_SCT = TypeVar("_SCT", bound=generic)
# The returned arrays dtype must be compatible with `np.equal`
_MaskFunc = Callable[
[NDArray[int_], _T],
NDArray[number[Any] | np.bool | timedelta64 | datetime64 | object_],
]
__all__: list[str]
@overload
def fliplr(m: _ArrayLike[_SCT]) -> NDArray[_SCT]: ...
@overload
def fliplr(m: ArrayLike) -> NDArray[Any]: ...
@overload
def flipud(m: _ArrayLike[_SCT]) -> NDArray[_SCT]: ...
@overload
def flipud(m: ArrayLike) -> NDArray[Any]: ...
@overload
def eye(
N: int,
M: None | int = ...,
k: int = ...,
dtype: None = ...,
order: _OrderCF = ...,
*,
device: None | L["cpu"] = ...,
like: None | _SupportsArrayFunc = ...,
) -> NDArray[float64]: ...
@overload
def eye(
N: int,
M: None | int = ...,
k: int = ...,
dtype: _DTypeLike[_SCT] = ...,
order: _OrderCF = ...,
*,
device: None | L["cpu"] = ...,
like: None | _SupportsArrayFunc = ...,
) -> NDArray[_SCT]: ...
@overload
def eye(
N: int,
M: None | int = ...,
k: int = ...,
dtype: DTypeLike = ...,
order: _OrderCF = ...,
*,
device: None | L["cpu"] = ...,
like: None | _SupportsArrayFunc = ...,
) -> NDArray[Any]: ...
@overload
def diag(v: _ArrayLike[_SCT], k: int = ...) -> NDArray[_SCT]: ...
@overload
def diag(v: ArrayLike, k: int = ...) -> NDArray[Any]: ...
@overload
def diagflat(v: _ArrayLike[_SCT], k: int = ...) -> NDArray[_SCT]: ...
@overload
def diagflat(v: ArrayLike, k: int = ...) -> NDArray[Any]: ...
@overload
def tri(
N: int,
M: None | int = ...,
k: int = ...,
dtype: None = ...,
*,
like: None | _SupportsArrayFunc = ...
) -> NDArray[float64]: ...
@overload
def tri(
N: int,
M: None | int = ...,
k: int = ...,
dtype: _DTypeLike[_SCT] = ...,
*,
like: None | _SupportsArrayFunc = ...
) -> NDArray[_SCT]: ...
@overload
def tri(
N: int,
M: None | int = ...,
k: int = ...,
dtype: DTypeLike = ...,
*,
like: None | _SupportsArrayFunc = ...
) -> NDArray[Any]: ...
@overload
def tril(v: _ArrayLike[_SCT], k: int = ...) -> NDArray[_SCT]: ...
@overload
def tril(v: ArrayLike, k: int = ...) -> NDArray[Any]: ...
@overload
def triu(v: _ArrayLike[_SCT], k: int = ...) -> NDArray[_SCT]: ...
@overload
def triu(v: ArrayLike, k: int = ...) -> NDArray[Any]: ...
@overload
def vander( # type: ignore[misc]
x: _ArrayLikeInt_co,
N: None | int = ...,
increasing: bool = ...,
) -> NDArray[signedinteger[Any]]: ...
@overload
def vander( # type: ignore[misc]
x: _ArrayLikeFloat_co,
N: None | int = ...,
increasing: bool = ...,
) -> NDArray[floating[Any]]: ...
@overload
def vander(
x: _ArrayLikeComplex_co,
N: None | int = ...,
increasing: bool = ...,
) -> NDArray[complexfloating[Any, Any]]: ...
@overload
def vander(
x: _ArrayLikeObject_co,
N: None | int = ...,
increasing: bool = ...,
) -> NDArray[object_]: ...
_Int_co: TypeAlias = np.integer[Any] | np.bool
_Float_co: TypeAlias = np.floating[Any] | _Int_co
_Number_co: TypeAlias = np.number[Any] | np.bool
_ArrayLike1D: TypeAlias = _SupportsArray[np.dtype[_SCT]] | Sequence[_SCT]
_ArrayLike2D: TypeAlias = (
_SupportsArray[np.dtype[_SCT]]
| Sequence[_ArrayLike1D[_SCT]]
)
_ArrayLike1DInt_co = (
_SupportsArray[np.dtype[_Int_co]]
| Sequence[int | _Int_co]
)
_ArrayLike1DFloat_co = (
_SupportsArray[np.dtype[_Float_co]]
| Sequence[float | int | _Float_co]
)
_ArrayLike2DFloat_co = (
_SupportsArray[np.dtype[_Float_co]]
| Sequence[_ArrayLike1DFloat_co]
)
_ArrayLike1DNumber_co = (
_SupportsArray[np.dtype[_Number_co]]
| Sequence[int | float | complex | _Number_co]
)
_SCT_complex = TypeVar("_SCT_complex", bound=np.complexfloating[Any, Any])
_SCT_inexact = TypeVar("_SCT_inexact", bound=np.inexact[Any])
_SCT_number_co = TypeVar("_SCT_number_co", bound=_Number_co)
@overload
def histogram2d(
x: _ArrayLike1D[_SCT_complex],
y: _ArrayLike1D[_SCT_complex | _Float_co],
bins: int | Sequence[int] = ...,
range: None | _ArrayLike2DFloat_co = ...,
density: None | bool = ...,
weights: None | _ArrayLike1DFloat_co = ...,
) -> tuple[
NDArray[float64],
NDArray[_SCT_complex],
NDArray[_SCT_complex],
]: ...
@overload
def histogram2d(
x: _ArrayLike1D[_SCT_complex | _Float_co],
y: _ArrayLike1D[_SCT_complex],
bins: int | Sequence[int] = ...,
range: None | _ArrayLike2DFloat_co = ...,
density: None | bool = ...,
weights: None | _ArrayLike1DFloat_co = ...,
) -> tuple[
NDArray[float64],
NDArray[_SCT_complex],
NDArray[_SCT_complex],
]: ...
@overload
def histogram2d(
x: _ArrayLike1D[_SCT_inexact],
y: _ArrayLike1D[_SCT_inexact | _Int_co],
bins: int | Sequence[int] = ...,
range: None | _ArrayLike2DFloat_co = ...,
density: None | bool = ...,
weights: None | _ArrayLike1DFloat_co = ...,
) -> tuple[
NDArray[float64],
NDArray[_SCT_inexact],
NDArray[_SCT_inexact],
]: ...
@overload
def histogram2d(
x: _ArrayLike1D[_SCT_inexact | _Int_co],
y: _ArrayLike1D[_SCT_inexact],
bins: int | Sequence[int] = ...,
range: None | _ArrayLike2DFloat_co = ...,
density: None | bool = ...,
weights: None | _ArrayLike1DFloat_co = ...,
) -> tuple[
NDArray[float64],
NDArray[_SCT_inexact],
NDArray[_SCT_inexact],
]: ...
@overload
def histogram2d(
x: _ArrayLike1DInt_co | Sequence[float | int],
y: _ArrayLike1DInt_co | Sequence[float | int],
bins: int | Sequence[int] = ...,
range: None | _ArrayLike2DFloat_co = ...,
density: None | bool = ...,
weights: None | _ArrayLike1DFloat_co = ...,
) -> tuple[
NDArray[float64],
NDArray[float64],
NDArray[float64],
]: ...
@overload
def histogram2d(
x: Sequence[complex | float | int],
y: Sequence[complex | float | int],
bins: int | Sequence[int] = ...,
range: None | _ArrayLike2DFloat_co = ...,
density: None | bool = ...,
weights: None | _ArrayLike1DFloat_co = ...,
) -> tuple[
NDArray[float64],
NDArray[complex128 | float64],
NDArray[complex128 | float64],
]: ...
@overload
def histogram2d(
x: _ArrayLike1DNumber_co,
y: _ArrayLike1DNumber_co,
bins: _ArrayLike1D[_SCT_number_co] | Sequence[_ArrayLike1D[_SCT_number_co]],
range: None | _ArrayLike2DFloat_co = ...,
density: None | bool = ...,
weights: None | _ArrayLike1DFloat_co = ...,
) -> tuple[
NDArray[float64],
NDArray[_SCT_number_co],
NDArray[_SCT_number_co],
]: ...
@overload
def histogram2d(
x: _ArrayLike1D[_SCT_inexact],
y: _ArrayLike1D[_SCT_inexact],
bins: Sequence[_ArrayLike1D[_SCT_number_co] | int],
range: None | _ArrayLike2DFloat_co = ...,
density: None | bool = ...,
weights: None | _ArrayLike1DFloat_co = ...,
) -> tuple[
NDArray[float64],
NDArray[_SCT_number_co | _SCT_inexact],
NDArray[_SCT_number_co | _SCT_inexact],
]: ...
@overload
def histogram2d(
x: _ArrayLike1DInt_co | Sequence[float | int],
y: _ArrayLike1DInt_co | Sequence[float | int],
bins: Sequence[_ArrayLike1D[_SCT_number_co] | int],
range: None | _ArrayLike2DFloat_co = ...,
density: None | bool = ...,
weights: None | _ArrayLike1DFloat_co = ...,
) -> tuple[
NDArray[float64],
NDArray[_SCT_number_co | float64],
NDArray[_SCT_number_co | float64],
]: ...
@overload
def histogram2d(
x: Sequence[complex | float | int],
y: Sequence[complex | float | int],
bins: Sequence[_ArrayLike1D[_SCT_number_co] | int],
range: None | _ArrayLike2DFloat_co = ...,
density: None | bool = ...,
weights: None | _ArrayLike1DFloat_co = ...,
) -> tuple[
NDArray[float64],
NDArray[_SCT_number_co | complex128 | float64],
NDArray[_SCT_number_co | complex128 | float64] ,
]: ...
@overload
def histogram2d(
x: _ArrayLike1DNumber_co,
y: _ArrayLike1DNumber_co,
bins: Sequence[Sequence[bool]],
range: None | _ArrayLike2DFloat_co = ...,
density: None | bool = ...,
weights: None | _ArrayLike1DFloat_co = ...,
) -> tuple[
NDArray[float64],
NDArray[np.bool],
NDArray[np.bool],
]: ...
@overload
def histogram2d(
x: _ArrayLike1DNumber_co,
y: _ArrayLike1DNumber_co,
bins: Sequence[Sequence[int | bool]],
range: None | _ArrayLike2DFloat_co = ...,
density: None | bool = ...,
weights: None | _ArrayLike1DFloat_co = ...,
) -> tuple[
NDArray[float64],
NDArray[np.int_ | np.bool],
NDArray[np.int_ | np.bool],
]: ...
@overload
def histogram2d(
x: _ArrayLike1DNumber_co,
y: _ArrayLike1DNumber_co,
bins: Sequence[Sequence[float | int | bool]],
range: None | _ArrayLike2DFloat_co = ...,
density: None | bool = ...,
weights: None | _ArrayLike1DFloat_co = ...,
) -> tuple[
NDArray[float64],
NDArray[np.float64 | np.int_ | np.bool],
NDArray[np.float64 | np.int_ | np.bool],
]: ...
@overload
def histogram2d(
x: _ArrayLike1DNumber_co,
y: _ArrayLike1DNumber_co,
bins: Sequence[Sequence[complex | float | int | bool]],
range: None | _ArrayLike2DFloat_co = ...,
density: None | bool = ...,
weights: None | _ArrayLike1DFloat_co = ...,
) -> tuple[
NDArray[float64],
NDArray[np.complex128 | np.float64 | np.int_ | np.bool],
NDArray[np.complex128 | np.float64 | np.int_ | np.bool],
]: ...
# NOTE: we're assuming/demanding here the `mask_func` returns
# an ndarray of shape `(n, n)`; otherwise there is the possibility
# of the output tuple having more or less than 2 elements
@overload
def mask_indices(
n: int,
mask_func: _MaskFunc[int],
k: int = ...,
) -> tuple[NDArray[intp], NDArray[intp]]: ...
@overload
def mask_indices(
n: int,
mask_func: _MaskFunc[_T],
k: _T,
) -> tuple[NDArray[intp], NDArray[intp]]: ...
def tril_indices(
n: int,
k: int = ...,
m: None | int = ...,
) -> tuple[NDArray[int_], NDArray[int_]]: ...
def tril_indices_from(
arr: NDArray[Any],
k: int = ...,
) -> tuple[NDArray[int_], NDArray[int_]]: ...
def triu_indices(
n: int,
k: int = ...,
m: None | int = ...,
) -> tuple[NDArray[int_], NDArray[int_]]: ...
def triu_indices_from(
arr: NDArray[Any],
k: int = ...,
) -> tuple[NDArray[int_], NDArray[int_]]: ...