| from collections.abc import Callable, Sequence
|
| from typing import (
|
| TypeVar,
|
| Any,
|
| overload,
|
| SupportsIndex,
|
| Protocol,
|
| ParamSpec,
|
| Concatenate,
|
| type_check_only,
|
| )
|
|
|
| import numpy as np
|
| from numpy import (
|
| generic,
|
| integer,
|
| ufunc,
|
| unsignedinteger,
|
| signedinteger,
|
| floating,
|
| complexfloating,
|
| object_,
|
| )
|
| from numpy._core.shape_base import vstack as row_stack
|
| from numpy._typing import (
|
| ArrayLike,
|
| NDArray,
|
| _ShapeLike,
|
| _ArrayLike,
|
| _ArrayLikeBool_co,
|
| _ArrayLikeUInt_co,
|
| _ArrayLikeInt_co,
|
| _ArrayLikeFloat_co,
|
| _ArrayLikeComplex_co,
|
| _ArrayLikeObject_co,
|
| )
|
|
|
| __all__ = [
|
| "column_stack",
|
| "row_stack",
|
| "dstack",
|
| "array_split",
|
| "split",
|
| "hsplit",
|
| "vsplit",
|
| "dsplit",
|
| "apply_over_axes",
|
| "expand_dims",
|
| "apply_along_axis",
|
| "kron",
|
| "tile",
|
| "take_along_axis",
|
| "put_along_axis",
|
| ]
|
|
|
| _P = ParamSpec("_P")
|
| _SCT = TypeVar("_SCT", bound=generic)
|
|
|
| # Signature of `__array_wrap__`
|
| @type_check_only
|
| class _ArrayWrap(Protocol):
|
| def __call__(
|
| self,
|
| array: NDArray[Any],
|
| context: None | tuple[ufunc, tuple[Any, ...], int] = ...,
|
| return_scalar: bool = ...,
|
| /,
|
| ) -> Any: ...
|
|
|
| @type_check_only
|
| class _SupportsArrayWrap(Protocol):
|
| @property
|
| def __array_wrap__(self) -> _ArrayWrap: ...
|
|
|
| def take_along_axis(
|
| arr: _SCT | NDArray[_SCT],
|
| indices: NDArray[integer[Any]],
|
| axis: None | int,
|
| ) -> NDArray[_SCT]: ...
|
|
|
| def put_along_axis(
|
| arr: NDArray[_SCT],
|
| indices: NDArray[integer[Any]],
|
| values: ArrayLike,
|
| axis: None | int,
|
| ) -> None: ...
|
|
|
| @overload
|
| def apply_along_axis(
|
| func1d: Callable[Concatenate[NDArray[Any], _P], _ArrayLike[_SCT]],
|
| axis: SupportsIndex,
|
| arr: ArrayLike,
|
| *args: _P.args,
|
| **kwargs: _P.kwargs,
|
| ) -> NDArray[_SCT]: ...
|
| @overload
|
| def apply_along_axis(
|
| func1d: Callable[Concatenate[NDArray[Any], _P], Any],
|
| axis: SupportsIndex,
|
| arr: ArrayLike,
|
| *args: _P.args,
|
| **kwargs: _P.kwargs,
|
| ) -> NDArray[Any]: ...
|
|
|
| def apply_over_axes(
|
| func: Callable[[NDArray[Any], int], NDArray[_SCT]],
|
| a: ArrayLike,
|
| axes: int | Sequence[int],
|
| ) -> NDArray[_SCT]: ...
|
|
|
| @overload
|
| def expand_dims(
|
| a: _ArrayLike[_SCT],
|
| axis: _ShapeLike,
|
| ) -> NDArray[_SCT]: ...
|
| @overload
|
| def expand_dims(
|
| a: ArrayLike,
|
| axis: _ShapeLike,
|
| ) -> NDArray[Any]: ...
|
|
|
| @overload
|
| def column_stack(tup: Sequence[_ArrayLike[_SCT]]) -> NDArray[_SCT]: ...
|
| @overload
|
| def column_stack(tup: Sequence[ArrayLike]) -> NDArray[Any]: ...
|
|
|
| @overload
|
| def dstack(tup: Sequence[_ArrayLike[_SCT]]) -> NDArray[_SCT]: ...
|
| @overload
|
| def dstack(tup: Sequence[ArrayLike]) -> NDArray[Any]: ...
|
|
|
| @overload
|
| def array_split(
|
| ary: _ArrayLike[_SCT],
|
| indices_or_sections: _ShapeLike,
|
| axis: SupportsIndex = ...,
|
| ) -> list[NDArray[_SCT]]: ...
|
| @overload
|
| def array_split(
|
| ary: ArrayLike,
|
| indices_or_sections: _ShapeLike,
|
| axis: SupportsIndex = ...,
|
| ) -> list[NDArray[Any]]: ...
|
|
|
| @overload
|
| def split(
|
| ary: _ArrayLike[_SCT],
|
| indices_or_sections: _ShapeLike,
|
| axis: SupportsIndex = ...,
|
| ) -> list[NDArray[_SCT]]: ...
|
| @overload
|
| def split(
|
| ary: ArrayLike,
|
| indices_or_sections: _ShapeLike,
|
| axis: SupportsIndex = ...,
|
| ) -> list[NDArray[Any]]: ...
|
|
|
| @overload
|
| def hsplit(
|
| ary: _ArrayLike[_SCT],
|
| indices_or_sections: _ShapeLike,
|
| ) -> list[NDArray[_SCT]]: ...
|
| @overload
|
| def hsplit(
|
| ary: ArrayLike,
|
| indices_or_sections: _ShapeLike,
|
| ) -> list[NDArray[Any]]: ...
|
|
|
| @overload
|
| def vsplit(
|
| ary: _ArrayLike[_SCT],
|
| indices_or_sections: _ShapeLike,
|
| ) -> list[NDArray[_SCT]]: ...
|
| @overload
|
| def vsplit(
|
| ary: ArrayLike,
|
| indices_or_sections: _ShapeLike,
|
| ) -> list[NDArray[Any]]: ...
|
|
|
| @overload
|
| def dsplit(
|
| ary: _ArrayLike[_SCT],
|
| indices_or_sections: _ShapeLike,
|
| ) -> list[NDArray[_SCT]]: ...
|
| @overload
|
| def dsplit(
|
| ary: ArrayLike,
|
| indices_or_sections: _ShapeLike,
|
| ) -> list[NDArray[Any]]: ...
|
|
|
| @overload
|
| def get_array_wrap(*args: _SupportsArrayWrap) -> _ArrayWrap: ...
|
| @overload
|
| def get_array_wrap(*args: object) -> None | _ArrayWrap: ...
|
|
|
| @overload
|
| def kron(a: _ArrayLikeBool_co, b: _ArrayLikeBool_co) -> NDArray[np.bool]: ... # type: ignore[misc]
|
| @overload
|
| def kron(a: _ArrayLikeUInt_co, b: _ArrayLikeUInt_co) -> NDArray[unsignedinteger[Any]]: ... # type: ignore[misc]
|
| @overload
|
| def kron(a: _ArrayLikeInt_co, b: _ArrayLikeInt_co) -> NDArray[signedinteger[Any]]: ... # type: ignore[misc]
|
| @overload
|
| def kron(a: _ArrayLikeFloat_co, b: _ArrayLikeFloat_co) -> NDArray[floating[Any]]: ... # type: ignore[misc]
|
| @overload
|
| def kron(a: _ArrayLikeComplex_co, b: _ArrayLikeComplex_co) -> NDArray[complexfloating[Any, Any]]: ...
|
| @overload
|
| def kron(a: _ArrayLikeObject_co, b: Any) -> NDArray[object_]: ...
|
| @overload
|
| def kron(a: Any, b: _ArrayLikeObject_co) -> NDArray[object_]: ...
|
|
|
| @overload
|
| def tile(
|
| A: _ArrayLike[_SCT],
|
| reps: int | Sequence[int],
|
| ) -> NDArray[_SCT]: ...
|
| @overload
|
| def tile(
|
| A: ArrayLike,
|
| reps: int | Sequence[int],
|
| ) -> NDArray[Any]: ...
|
| |
