import sys from abc import abstractmethod from dataclasses import dataclass from typing import ( Any, Callable, Container, Dict, Generic, Iterable, Iterator, List, Optional, Sequence, Sized, Tuple, TypeVar, Union, cast, overload, ) import numpy from .compat import cupy, has_cupy if has_cupy: get_array_module = cupy.get_array_module else: get_array_module = lambda obj: numpy # Use typing_extensions for Python versions < 3.8 if sys.version_info < (3, 8): from typing_extensions import Literal, Protocol else: from typing import Literal, Protocol # noqa: F401 # fmt: off XY_YZ_OutT = TypeVar("XY_YZ_OutT") XY_XY_OutT = TypeVar("XY_XY_OutT") DeviceTypes = Literal["cpu", "gpu", "tpu"] Batchable = Union["Pairs", "Ragged", "Padded", "ArrayXd", List, Tuple] Xp = Union["numpy", "cupy"] # type: ignore Shape = Tuple[int, ...] DTypes = Literal["f", "i", "float16", "float32", "float64", "int32", "int64", "uint32", "uint64"] DTypesFloat = Literal["f", "float32", "float16", "float64"] DTypesInt = Literal["i", "int32", "int64", "uint32", "uint64"] Array1d = Union["Floats1d", "Ints1d"] Array2d = Union["Floats2d", "Ints2d"] Array3d = Union["Floats3d", "Ints3d"] Array4d = Union["Floats4d", "Ints4d"] FloatsXd = Union["Floats1d", "Floats2d", "Floats3d", "Floats4d"] IntsXd = Union["Ints1d", "Ints2d", "Ints3d", "Ints4d"] ArrayXd = Union[FloatsXd, IntsXd] List1d = Union[List["Floats1d"], List["Ints1d"]] List2d = Union[List["Floats2d"], List["Ints2d"]] List3d = Union[List["Floats3d"], List["Ints3d"]] List4d = Union[List["Floats4d"], List["Ints4d"]] ListXd = Union[List1d, List2d, List3d, List4d] ArrayT = TypeVar("ArrayT") SelfT = TypeVar("SelfT") Array1dT = TypeVar("Array1dT", bound="Array1d") FloatsXdT = TypeVar("FloatsXdT", "Floats1d", "Floats2d", "Floats3d", "Floats4d") # These all behave the same as far as indexing is concerned Slicish = Union[slice, List[int], "ArrayXd"] _1_KeyScalar = int _1_Key1d = Slicish _1_AllKeys = Union[_1_KeyScalar, _1_Key1d] _F1_AllReturns = Union[float, "Floats1d"] _I1_AllReturns = Union[int, "Ints1d"] _2_KeyScalar = Tuple[int, int] _2_Key1d = Union[int, Tuple[Slicish, int], Tuple[int, Slicish]] _2_Key2d = Union[Tuple[Slicish, Slicish], Slicish] _2_AllKeys = Union[_2_KeyScalar, _2_Key1d, _2_Key2d] _F2_AllReturns = Union[float, "Floats1d", "Floats2d"] _I2_AllReturns = Union[int, "Ints1d", "Ints2d"] _3_KeyScalar = Tuple[int, int, int] _3_Key1d = Union[Tuple[int, int], Tuple[int, int, Slicish], Tuple[int, Slicish, int], Tuple[Slicish, int, int]] _3_Key2d = Union[int, Tuple[int, Slicish], Tuple[Slicish, int], Tuple[int, Slicish, Slicish], Tuple[Slicish, int, Slicish], Tuple[Slicish, Slicish, int]] _3_Key3d = Union[Slicish, Tuple[Slicish, Slicish], Tuple[Slicish, Slicish, Slicish]] _3_AllKeys = Union[_3_KeyScalar, _3_Key1d, _3_Key2d, _3_Key3d] _F3_AllReturns = Union[float, "Floats1d", "Floats2d", "Floats3d"] _I3_AllReturns = Union[int, "Ints1d", "Ints2d", "Ints3d"] _4_KeyScalar = Tuple[int, int, int, int] _4_Key1d = Union[Tuple[int, int, int], Tuple[int, int, int, Slicish], Tuple[int, int, Slicish, int], Tuple[int, Slicish, int, int], Tuple[Slicish, int, int, int]] _4_Key2d = Union[Tuple[int, int], Tuple[int, int, Slicish], Tuple[int, Slicish, int], Tuple[Slicish, int, int], Tuple[int, int, Slicish, Slicish], Tuple[int, Slicish, int, Slicish], Tuple[int, Slicish, Slicish, int], Tuple[Slicish, int, int, Slicish], Tuple[Slicish, int, Slicish, int], Tuple[Slicish, Slicish, int, int]] _4_Key3d = Union[int, Tuple[int, Slicish], Tuple[Slicish, int], Tuple[int, Slicish, Slicish], Tuple[Slicish, int, Slicish], Tuple[Slicish, Slicish, int], Tuple[int, Slicish, Slicish, Slicish], Tuple[Slicish, int, Slicish, Slicish], Tuple[Slicish, Slicish, int, Slicish], Tuple[Slicish, Slicish, Slicish, int]] _4_Key4d = Union[Slicish, Tuple[Slicish, Slicish], Tuple[Slicish, Slicish, Slicish], Tuple[Slicish, Slicish, Slicish, Slicish]] _4_AllKeys = Union[_4_KeyScalar, _4_Key1d, _4_Key2d, _4_Key3d, _4_Key4d] _F4_AllReturns = Union[float, "Floats1d", "Floats2d", "Floats3d", "Floats4d"] _I4_AllReturns = Union[int, "Ints1d", "Ints2d", "Ints3d", "Ints4d"] # Typedefs for the reduction methods. Tru = Literal[True] Fal = Literal[False] OneAx = Union[int, Tuple[int]] TwoAx = Tuple[int, int] ThreeAx = Tuple[int, int, int] FourAx = Tuple[int, int, int, int] _1_AllAx = Optional[OneAx] _2_AllAx = Union[Optional[TwoAx], OneAx] _3_AllAx = Union[Optional[ThreeAx], TwoAx, OneAx] _4_AllAx = Union[Optional[FourAx], ThreeAx, TwoAx, OneAx] _1F_ReduceResults = Union[float, "Floats1d"] _2F_ReduceResults = Union[float, "Floats1d", "Floats2d"] _3F_ReduceResults = Union[float, "Floats1d", "Floats2d", "Floats3d"] _4F_ReduceResults = Union[float, "Floats1d", "Floats2d", "Floats3d", "Floats4d"] _1I_ReduceResults = Union[int, "Ints1d"] _2I_ReduceResults = Union[int, "Ints1d", "Ints2d"] _3I_ReduceResults = Union[int, "Ints1d", "Ints2d", "Ints3d"] _4I_ReduceResults = Union[int, "Ints1d", "Ints2d", "Ints3d", "Ints4d"] # TODO: # We need to get correct overloads in for the following reduction methods. # The 'sum' reduction is correct --- the others need to be just the same, # but with a different name. # max, min, prod, round, var, mean, ptp, std # There's also one *slightly* different function, cumsum. This doesn't # have a scalar version -- it always makes an array. class _Array(Sized, Container): @classmethod def __get_validators__(cls): """Runtime validation for pydantic.""" yield lambda v: validate_array(v) @property @abstractmethod def dtype(self) -> DTypes: ... @property @abstractmethod def data(self) -> memoryview: ... @property @abstractmethod def flags(self) -> Any: ... @property @abstractmethod def size(self) -> int: ... @property @abstractmethod def itemsize(self) -> int: ... @property @abstractmethod def nbytes(self) -> int: ... @property @abstractmethod def ndim(self) -> int: ... @property @abstractmethod def shape(self) -> Shape: ... @property @abstractmethod def strides(self) -> Tuple[int, ...]: ... # TODO: Is ArrayT right? @abstractmethod def astype(self: ArrayT, dtype: DTypes, order: str = ..., casting: str = ..., subok: bool = ..., copy: bool = ...) -> ArrayT: ... @abstractmethod def copy(self: ArrayT, order: str = ...) -> ArrayT: ... @abstractmethod def fill(self, value: Any) -> None: ... # Shape manipulation @abstractmethod def reshape(self: ArrayT, shape: Shape, *, order: str = ...) -> ArrayT: ... @abstractmethod def transpose(self: ArrayT, axes: Shape) -> ArrayT: ... # TODO: is this right? It returns 1d @abstractmethod def flatten(self, order: str = ...): ... # TODO: is this right? It returns 1d @abstractmethod def ravel(self, order: str = ...): ... @abstractmethod def squeeze(self, axis: Union[int, Shape] = ...): ... @abstractmethod def __len__(self) -> int: ... @abstractmethod def __setitem__(self, key, value): ... @abstractmethod def __iter__(self) -> Iterator[Any]: ... @abstractmethod def __contains__(self, key) -> bool: ... @abstractmethod def __index__(self) -> int: ... @abstractmethod def __int__(self) -> int: ... @abstractmethod def __float__(self) -> float: ... @abstractmethod def __complex__(self) -> complex: ... @abstractmethod def __bool__(self) -> bool: ... @abstractmethod def __bytes__(self) -> bytes: ... @abstractmethod def __str__(self) -> str: ... @abstractmethod def __repr__(self) -> str: ... @abstractmethod def __copy__(self, order: str = ...): ... @abstractmethod def __deepcopy__(self: SelfT, memo: dict) -> SelfT: ... @abstractmethod def __lt__(self, other): ... @abstractmethod def __le__(self, other): ... @abstractmethod def __eq__(self, other): ... @abstractmethod def __ne__(self, other): ... @abstractmethod def __gt__(self, other): ... @abstractmethod def __ge__(self, other): ... @abstractmethod def __add__(self, other): ... @abstractmethod def __radd__(self, other): ... @abstractmethod def __iadd__(self, other): ... @abstractmethod def __sub__(self, other): ... @abstractmethod def __rsub__(self, other): ... @abstractmethod def __isub__(self, other): ... @abstractmethod def __mul__(self, other): ... @abstractmethod def __rmul__(self, other): ... @abstractmethod def __imul__(self, other): ... @abstractmethod def __truediv__(self, other): ... @abstractmethod def __rtruediv__(self, other): ... @abstractmethod def __itruediv__(self, other): ... @abstractmethod def __floordiv__(self, other): ... @abstractmethod def __rfloordiv__(self, other): ... @abstractmethod def __ifloordiv__(self, other): ... @abstractmethod def __mod__(self, other): ... @abstractmethod def __rmod__(self, other): ... @abstractmethod def __imod__(self, other): ... @abstractmethod def __divmod__(self, other): ... @abstractmethod def __rdivmod__(self, other): ... # NumPy's __pow__ doesn't handle a third argument @abstractmethod def __pow__(self, other): ... @abstractmethod def __rpow__(self, other): ... @abstractmethod def __ipow__(self, other): ... @abstractmethod def __lshift__(self, other): ... @abstractmethod def __rlshift__(self, other): ... @abstractmethod def __ilshift__(self, other): ... @abstractmethod def __rshift__(self, other): ... @abstractmethod def __rrshift__(self, other): ... @abstractmethod def __irshift__(self, other): ... @abstractmethod def __and__(self, other): ... @abstractmethod def __rand__(self, other): ... @abstractmethod def __iand__(self, other): ... @abstractmethod def __xor__(self, other): ... @abstractmethod def __rxor__(self, other): ... @abstractmethod def __ixor__(self, other): ... @abstractmethod def __or__(self, other): ... @abstractmethod def __ror__(self, other): ... @abstractmethod def __ior__(self, other): ... @abstractmethod def __matmul__(self, other): ... @abstractmethod def __rmatmul__(self, other): ... @abstractmethod def __neg__(self: ArrayT) -> ArrayT: ... @abstractmethod def __pos__(self: ArrayT) -> ArrayT: ... @abstractmethod def __abs__(self: ArrayT) -> ArrayT: ... @abstractmethod def __invert__(self: ArrayT) -> ArrayT: ... @abstractmethod def get(self: ArrayT) -> ArrayT: ... @abstractmethod def all(self, axis: int = -1, out: Optional[ArrayT] = None, keepdims: bool = False) -> ArrayT: ... @abstractmethod def any(self, axis: int = -1, out: Optional[ArrayT] = None, keepdims: bool = False) -> ArrayT: ... # def argmax(self, axis: int = -1, out: Optional["Array"] = None, keepdims: Union[Tru, Fal]=False) -> Union[int, "Ints1d"]: ... @abstractmethod def argmin(self, axis: int = -1, out: Optional[ArrayT] = None) -> ArrayT: ... @abstractmethod def clip(self, a_min: Any, a_max: Any, out: Optional[ArrayT]) -> ArrayT: ... #def cumsum( self: ArrayT, axis: int = -1, dtype: Optional[DTypes] = None, out: Optional[ArrayT] = None) -> ArrayT: ... @abstractmethod def max(self, axis: int = -1, out: Optional[ArrayT] = None) -> ArrayT: ... # def mean(self, axis: int = -1, dtype: Optional[DTypes] = None, out: Optional[SelfT] = None, keepdims: bool = False) -> "Array": ... @abstractmethod def min(self, axis: int = -1, out: Optional[ArrayT] = None) -> ArrayT: ... @abstractmethod def nonzero(self: SelfT) -> SelfT: ... @abstractmethod def prod(self, axis: int = -1, dtype: Optional[DTypes] = None, out: Optional[ArrayT] = None, keepdims: bool = False) -> ArrayT: ... @abstractmethod def round(self, decimals: int = 0, out: Optional[ArrayT] = None) -> ArrayT: ... # def sum(self, axis: int = -1, dtype: Optional[DTypes] = None, out: Optional[ArrayT] = None, keepdims: bool = False) -> ArrayT: ... @abstractmethod def tobytes(self, order: str = "C") -> bytes: ... @abstractmethod def tolist(self) -> List[Any]: ... @abstractmethod def var(self: SelfT, axis: int = -1, dtype: Optional[DTypes] = None, out: Optional[ArrayT] = None, ddof: int = 0, keepdims: bool = False) -> SelfT: ... class _Floats(_Array): @property @abstractmethod def dtype(self) -> DTypesFloat: ... @abstractmethod def fill(self, value: float) -> None: ... @abstractmethod def reshape(self, shape: Shape, *, order: str = ...) -> "_Floats": ... class _Ints(_Array): @property @abstractmethod def dtype(self) -> DTypesInt: ... @abstractmethod def fill(self, value: int) -> None: ... @abstractmethod def reshape(self, shape: Shape, *, order: str = ...) -> "_Ints": ... """ Extensive overloads to represent __getitem__ behaviour. In an N+1 dimensional array, there will be N possible return types. For instance, if you have a 2d array, you could get back a float (array[i, j]), a floats1d (array[i]) or a floats2d (array[:i, :j]). You'll get the scalar if you have N ints in the index, a 1d array if you have N-1 ints, etc. So the trick here is to make a union with the various combinations that produce each result type, and then only have one overload per result. If we overloaded on each *key* type, that would get crazy, because there's tonnes of combinations. In each rank, we can use the same key-types for float and int, but we need a different return-type union. """ class _Array1d(_Array): """1-dimensional array.""" @classmethod def __get_validators__(cls): """Runtime validation for pydantic.""" yield lambda v: validate_array(v, ndim=1) @property @abstractmethod def ndim(self) -> Literal[1]: ... @property @abstractmethod def shape(self) -> Tuple[int]: ... @abstractmethod def __iter__(self) -> Iterator[Union[float, int]]: ... @abstractmethod def astype(self, dtype: DTypes, order: str = ..., casting: str = ..., subok: bool = ..., copy: bool = ...) -> "_Array1d": ... @abstractmethod def flatten(self: SelfT, order: str = ...) -> SelfT: ... @abstractmethod def ravel(self: SelfT, order: str = ...) -> SelfT: ... # These is actually a bit too strict: It's legal to say 'array1d + array2d' # That's kind of bad code though; it's better to write array2d + array1d. # We could relax this, but let's try the strict version. @abstractmethod def __add__(self: SelfT, other: Union[float, int, "Array1d"]) -> SelfT: ... @abstractmethod def __sub__(self: SelfT, other: Union[float, int, "Array1d"]) -> SelfT: ... @abstractmethod def __mul__(self: SelfT, other: Union[float, int, "Array1d"]) -> SelfT: ... @abstractmethod def __pow__(self: SelfT, other: Union[float, int, "Array1d"]) -> SelfT: ... @abstractmethod def __matmul__(self: SelfT, other: Union[float, int, "Array1d"]) -> SelfT: ... # These are not too strict though: you can't do += with higher dimensional. @abstractmethod def __iadd__(self, other: Union[float, int, "Array1d"]): ... @abstractmethod def __isub__(self, other: Union[float, int, "Array1d"]): ... @abstractmethod def __imul__(self, other: Union[float, int, "Array1d"]): ... @abstractmethod def __ipow__(self, other: Union[float, int, "Array1d"]): ... @overload @abstractmethod def argmax(self, keepdims: Fal = False, axis: int = -1, out: Optional[_Array] = None) -> int: ... @overload @abstractmethod def argmax(self, keepdims: Tru, axis: int = -1, out: Optional[_Array] = None) -> "Ints1d": ... @abstractmethod def argmax(self, keepdims: bool = False, axis: int = -1, out: Optional[_Array] = None) -> Union[int, "Ints1d"]: ... @overload @abstractmethod def mean(self, keepdims: Tru, axis: int = -1, dtype: Optional[DTypes] = None, out: Optional["Floats1d"] = None) -> "Floats1d": ... @overload @abstractmethod def mean(self, keepdims: Fal = False, axis: int = -1, dtype: Optional[DTypes] = None, out: Optional["Floats1d"] = None) -> float: ... @abstractmethod def mean(self, keepdims: bool = False, axis: int = -1, dtype: Optional[DTypes] = None, out: Optional["Floats1d"] = None) -> Union["Floats1d", float]: ... class Floats1d(_Array1d, _Floats): """1-dimensional array of floats.""" T: "Floats1d" @classmethod def __get_validators__(cls): """Runtime validation for pydantic.""" yield lambda v: validate_array(v, ndim=1, dtype="f") @abstractmethod def __iter__(self) -> Iterator[float]: ... @overload @abstractmethod def __getitem__(self, key: _1_KeyScalar) -> float: ... @overload @abstractmethod def __getitem__(self, key: _1_Key1d) -> "Floats1d": ... @abstractmethod def __getitem__(self, key: _1_AllKeys) -> _F1_AllReturns: ... @overload @abstractmethod def __setitem__(self, key: _1_KeyScalar, value: float) -> None: ... @overload @abstractmethod def __setitem__(self, key: _1_Key1d, value: "Floats1d") -> None: ... @abstractmethod def __setitem__(self, key: _1_AllKeys, _F1_AllReturns) -> None: ... @overload @abstractmethod def cumsum(self, *, keepdims: Tru, axis: Optional[OneAx] = None, out: Optional["Floats1d"] = None) -> "Floats1d": ... @overload # Cumsum is unusual in this @abstractmethod def cumsum(self, *, keepdims: Fal, axis: Optional[OneAx] = None, out: Optional["Floats1d"] = None) -> "Floats1d": ... @abstractmethod def cumsum(self, *, keepdims: bool = False, axis: _1_AllAx = None, out: Optional["Floats1d"] = None) -> "Floats1d": ... @overload @abstractmethod def sum(self, *, keepdims: Tru, axis: Optional[OneAx] = None, out: Optional["Floats1d"] = None) -> "Floats1d": ... @overload @abstractmethod def sum(self, *, keepdims: Fal, axis: Optional[OneAx] = None, out = None) -> float: ... @abstractmethod def sum(self, *, keepdims: bool = False, axis: _1_AllAx = None, out: Optional["Floats1d"] = None) -> _1F_ReduceResults: ... class Ints1d(_Array1d, _Ints): """1-dimensional array of ints.""" T: "Ints1d" @classmethod def __get_validators__(cls): """Runtime validation for pydantic.""" yield lambda v: validate_array(v, ndim=1, dtype="i") @abstractmethod def __iter__(self) -> Iterator[int]: ... @overload @abstractmethod def __getitem__(self, key: _1_KeyScalar) -> int: ... @overload @abstractmethod def __getitem__(self, key: _1_Key1d) -> "Ints1d": ... @abstractmethod def __getitem__(self, key: _1_AllKeys) -> _I1_AllReturns: ... @overload @abstractmethod def __setitem__(self, key: _1_KeyScalar, value: int) -> None: ... @overload @abstractmethod def __setitem__(self, key: _1_Key1d, value: Union[int, "Ints1d"]) -> None: ... @abstractmethod def __setitem__(self, key: _1_AllKeys, _I1_AllReturns) -> None: ... @overload @abstractmethod def cumsum(self, *, keepdims: Tru, axis: Optional[OneAx] = None, out: Optional["Ints1d"] = None) -> "Ints1d": ... @overload @abstractmethod def cumsum(self, *, keepdims: Fal = False, axis: Optional[OneAx] = None, out: Optional["Ints1d"] = None) -> "Ints1d": ... @abstractmethod def cumsum(self, *, keepdims: bool = False, axis: _1_AllAx = None, out: Optional["Ints1d"] = None) -> "Ints1d": ... @overload @abstractmethod def sum(self, *, keepdims: Tru, axis: Optional[OneAx] = None, out: Optional["Ints1d"] = None) -> "Ints1d": ... @overload @abstractmethod def sum(self, *, keepdims: Fal = False, axis: Optional[OneAx] = None, out = None) -> int: ... @abstractmethod def sum(self, *, keepdims: bool = False, axis: _1_AllAx = None, out: Optional["Ints1d"] = None) -> _1I_ReduceResults: ... class _Array2d(_Array): @classmethod def __get_validators__(cls): """Runtime validation for pydantic.""" yield lambda v: validate_array(v, ndim=2) @property @abstractmethod def ndim(self) -> Literal[2]: ... @property @abstractmethod def shape(self) -> Tuple[int, int]: ... @abstractmethod def __iter__(self) -> Iterator[Array1d]: ... @abstractmethod def astype(self, dtype: DTypes, order: str = ..., casting: str = ..., subok: bool = ..., copy: bool = ...) -> "Array2d": ... # These is actually a bit too strict: It's legal to say 'array2d + array3d' # That's kind of bad code though; it's better to write array3d + array2d. # We could relax this, but let's try the strict version. @abstractmethod def __add__(self: ArrayT, other: Union[float, int, Array1d, "Array2d"]) -> ArrayT: ... @abstractmethod def __sub__(self: ArrayT, other: Union[float, int, Array1d, "Array2d"]) -> ArrayT: ... @abstractmethod def __mul__(self: ArrayT, other: Union[float, int, Array1d, "Array2d"]) -> ArrayT: ... @abstractmethod def __pow__(self: ArrayT, other: Union[float, int, Array1d, "Array2d"]) -> ArrayT: ... @abstractmethod def __matmul__(self: ArrayT, other: Union[float, int, Array1d, "Array2d"]) -> ArrayT: ... # These are not too strict though: you can't do += with higher dimensional. @abstractmethod def __iadd__(self, other: Union[float, int, Array1d, "Array2d"]): ... @abstractmethod def __isub__(self, other: Union[float, int, Array1d, "Array2d"]): ... @abstractmethod def __imul__(self, other: Union[float, int, Array1d, "Array2d"]): ... @abstractmethod def __ipow__(self, other: Union[float, int, Array1d, "Array2d"]): ... @overload @abstractmethod def argmax(self, keepdims: Fal = False, axis: int = -1, out: Optional[_Array] = None) -> Ints1d: ... @overload @abstractmethod def argmax(self, keepdims: Tru, axis: int = -1, out: Optional[_Array] = None) -> "Ints2d": ... @abstractmethod def argmax(self, keepdims: bool = False, axis: int = -1, out: Optional[_Array] = None) -> Union[Ints1d, "Ints2d"]: ... @overload @abstractmethod def mean(self, keepdims: Fal = False, axis: int = -1, dtype: Optional[DTypes] = None, out: Optional["Floats2d"] = None) -> Floats1d: ... @overload @abstractmethod def mean(self, keepdims: Tru, axis: int = -1, dtype: Optional[DTypes] = None, out: Optional["Floats2d"] = None) -> "Floats2d": ... @abstractmethod def mean(self, keepdims: bool = False, axis: int = -1, dtype: Optional[DTypes] = None, out: Optional["Floats2d"] = None) -> Union["Floats2d", Floats1d]: ... class Floats2d(_Array2d, _Floats): """2-dimensional array of floats""" T: "Floats2d" @classmethod def __get_validators__(cls): """Runtime validation for pydantic.""" yield lambda v: validate_array(v, ndim=2, dtype="f") @abstractmethod def __iter__(self) -> Iterator[Floats1d]: ... @overload @abstractmethod def __getitem__(self, key: _2_KeyScalar) -> float: ... @overload @abstractmethod def __getitem__(self, key: _2_Key1d) -> Floats1d: ... @overload @abstractmethod def __getitem__(self, key: _2_Key2d) -> "Floats2d": ... @abstractmethod def __getitem__(self, key: _2_AllKeys) -> _F2_AllReturns: ... @overload @abstractmethod def __setitem__(self, key: _2_KeyScalar, value: float) -> None: ... @overload @abstractmethod def __setitem__(self, key: _2_Key1d, value: Union[float, Floats1d]) -> None: ... @overload @abstractmethod def __setitem__(self, key: _2_Key2d, value: _F2_AllReturns) -> None: ... @abstractmethod def __setitem__(self, key: _2_AllKeys, value: _F2_AllReturns) -> None: ... @overload @abstractmethod def sum(self, *, keepdims: Tru, axis: _2_AllAx = None, out: Optional["Floats2d"] = None) -> "Floats2d": ... @overload @abstractmethod def sum(self, *, keepdims: Fal = False, axis: OneAx, out: Optional[Floats1d] = None) -> Floats1d: ... @overload @abstractmethod def sum(self, *, keepdims: Fal = False, axis: TwoAx, out = None) -> float: ... @abstractmethod def sum(self, *, keepdims: bool = False, axis: _2_AllAx = None, out: Union[None, "Floats1d", "Floats2d"] = None) -> _2F_ReduceResults: ... class Ints2d(_Array2d, _Ints): """2-dimensional array of ints.""" T: "Ints2d" @classmethod def __get_validators__(cls): """Runtime validation for pydantic.""" yield lambda v: validate_array(v, ndim=2, dtype="i") @abstractmethod def __iter__(self) -> Iterator[Ints1d]: ... @overload @abstractmethod def __getitem__(self, key: _2_KeyScalar) -> int: ... @overload @abstractmethod def __getitem__(self, key: _2_Key1d) -> Ints1d: ... @overload @abstractmethod def __getitem__(self, key: _2_Key2d) -> "Ints2d": ... @abstractmethod def __getitem__(self, key: _2_AllKeys) -> _I2_AllReturns: ... @overload @abstractmethod def __setitem__(self, key: _2_KeyScalar, value: int) -> None: ... @overload @abstractmethod def __setitem__(self, key: _2_Key1d, value: Ints1d) -> None: ... @overload @abstractmethod def __setitem__(self, key: _2_Key2d, value: "Ints2d") -> None: ... @abstractmethod def __setitem__(self, key: _2_AllKeys, value: _I2_AllReturns) -> None: ... @overload @abstractmethod def sum(self, keepdims: Fal = False, axis: int = -1, out: Optional["Ints1d"] = None) -> Ints1d: ... @overload @abstractmethod def sum(self, keepdims: Tru, axis: int = -1, out: Optional["Ints2d"] = None) -> "Ints2d": ... @abstractmethod def sum(self, keepdims: bool = False, axis: int = -1, out: Optional[Union["Ints1d", "Ints2d"]] = None) -> Union["Ints2d", Ints1d]: ... class _Array3d(_Array): """3-dimensional array of floats""" @classmethod def __get_validators__(cls): """Runtime validation for pydantic.""" yield lambda v: validate_array(v, ndim=3) @property @abstractmethod def ndim(self) -> Literal[3]: ... @property @abstractmethod def shape(self) -> Tuple[int, int, int]: ... @abstractmethod def __iter__(self) -> Iterator[Array2d]: ... @abstractmethod def astype(self, dtype: DTypes, order: str = ..., casting: str = ..., subok: bool = ..., copy: bool = ...) -> "Array3d": ... # These is actually a bit too strict: It's legal to say 'array2d + array3d' # That's kind of bad code though; it's better to write array3d + array2d. # We could relax this, but let's try the strict version. @abstractmethod def __add__(self: SelfT, other: Union[float, int, Array1d, Array2d, "Array3d"]) -> SelfT: ... @abstractmethod def __sub__(self: SelfT, other: Union[float, int, Array1d, Array2d, "Array3d"]) -> SelfT: ... @abstractmethod def __mul__(self: SelfT, other: Union[float, int, Array1d, Array2d, "Array3d"]) -> SelfT: ... @abstractmethod def __pow__(self: SelfT, other: Union[float, int, Array1d, Array2d, "Array3d"]) -> SelfT: ... @abstractmethod def __matmul__(self: SelfT, other: Union[float, int, Array1d, Array2d, "Array3d"]) -> SelfT: ... # These are not too strict though: you can't do += with higher dimensional. @abstractmethod def __iadd__(self, other: Union[float, int, Array1d, Array2d, "Array3d"]): ... @abstractmethod def __isub__(self, other: Union[float, int, Array1d, Array2d, "Array3d"]): ... @abstractmethod def __imul__(self, other: Union[float, int, Array1d, Array2d, "Array3d"]): ... @abstractmethod def __ipow__(self, other: Union[float, int, Array1d, Array2d, "Array3d"]): ... @overload @abstractmethod def argmax(self, keepdims: Fal = False, axis: int = -1, out: Optional[_Array] = None) -> Ints2d: ... @overload @abstractmethod def argmax(self, keepdims: Tru, axis: int = -1, out: Optional[_Array] = None) -> "Ints3d": ... @abstractmethod def argmax(self, keepdims: bool = False, axis: int = -1, out: Optional[_Array] = None) -> Union[Ints2d, "Ints3d"]: ... class Floats3d(_Array3d, _Floats): """3-dimensional array of floats""" T: "Floats3d" @classmethod def __get_validators__(cls): """Runtime validation for pydantic.""" yield lambda v: validate_array(v, ndim=3, dtype="f") @abstractmethod def __iter__(self) -> Iterator[Floats2d]: ... @overload @abstractmethod def __getitem__(self, key: _3_KeyScalar) -> float: ... @overload @abstractmethod def __getitem__(self, key: _3_Key1d) -> Floats1d: ... @overload @abstractmethod def __getitem__(self, key: _3_Key2d) -> Floats2d: ... @overload @abstractmethod def __getitem__(self, key: _3_Key3d) -> "Floats3d": ... @abstractmethod def __getitem__(self, key: _3_AllKeys) -> _F3_AllReturns: ... @overload @abstractmethod def __setitem__(self, key: _3_KeyScalar, value: float) -> None: ... @overload @abstractmethod def __setitem__(self, key: _3_Key1d, value: Floats1d) -> None: ... @overload @abstractmethod def __setitem__(self, key: _3_Key2d, value: Floats2d) -> None: ... @overload @abstractmethod def __setitem__(self, key: _3_Key3d, value: "Floats3d") -> None: ... @abstractmethod def __setitem__(self, key: _3_AllKeys, value: _F3_AllReturns) -> None: ... @overload @abstractmethod def sum(self, *, keepdims: Tru, axis: _3_AllAx = None, out: Optional["Floats3d"] = None) -> "Floats3d": ... @overload @abstractmethod def sum(self, *, keepdims: Fal, axis: OneAx, out: Optional[Floats2d] = None) -> Floats2d: ... @overload @abstractmethod def sum(self, *, keepdims: Fal, axis: TwoAx, out: Optional[Floats1d] = None) -> Floats1d: ... @overload @abstractmethod def sum(self, *, keepdims: Fal, axis: Optional[ThreeAx], out = None) -> float: ... @abstractmethod def sum(self, *, keepdims: bool = False, axis: _3_AllAx = None, out: Union[None, Floats1d, Floats2d, "Floats3d"] = None) -> _3F_ReduceResults: ... class Ints3d(_Array3d, _Ints): """3-dimensional array of ints.""" T: "Ints3d" @classmethod def __get_validators__(cls): """Runtime validation for pydantic.""" yield lambda v: validate_array(v, ndim=3, dtype="i") @abstractmethod def __iter__(self) -> Iterator[Ints2d]: ... @overload @abstractmethod def __getitem__(self, key: _3_KeyScalar) -> int: ... @overload @abstractmethod def __getitem__(self, key: _3_Key1d) -> Ints1d: ... @overload @abstractmethod def __getitem__(self, key: _3_Key2d) -> Ints2d: ... @overload @abstractmethod def __getitem__(self, key: _3_Key3d) -> "Ints3d": ... @abstractmethod def __getitem__(self, key: _3_AllKeys) -> _I3_AllReturns: ... @overload @abstractmethod def __setitem__(self, key: _3_KeyScalar, value: int) -> None: ... @overload @abstractmethod def __setitem__(self, key: _3_Key1d, value: Ints1d) -> None: ... @overload @abstractmethod def __setitem__(self, key: _3_Key2d, value: Ints2d) -> None: ... @overload @abstractmethod def __setitem__(self, key: _3_Key3d, value: "Ints3d") -> None: ... @abstractmethod def __setitem__(self, key: _3_AllKeys, value: _I3_AllReturns) -> None: ... @overload @abstractmethod def sum(self, *, keepdims: Tru, axis: _3_AllAx = None, out: Optional["Ints3d"] = None) -> "Ints3d": ... @overload @abstractmethod def sum(self, *, keepdims: Fal, axis: OneAx, out: Optional[Ints2d] = None) -> Ints2d: ... @overload @abstractmethod def sum(self, *, keepdims: Fal, axis: TwoAx, out: Optional[Ints1d] = None) -> Ints1d: ... @overload @abstractmethod def sum(self, *, keepdims: Fal, axis: Optional[ThreeAx], out = None) -> int: ... @abstractmethod def sum(self, *, keepdims: bool = False, axis: _3_AllAx = None, out: Union[None, Ints1d, Ints2d, "Ints3d"] = None) -> _3I_ReduceResults: ... class _Array4d(_Array): """4-dimensional array.""" @classmethod def __get_validators__(cls): """Runtime validation for pydantic.""" yield lambda v: validate_array(v, ndim=4) @property @abstractmethod def ndim(self) -> Literal[4]: ... @property @abstractmethod def shape(self) -> Tuple[int, int, int, int]: ... @abstractmethod def __iter__(self) -> Iterator[Array3d]: ... @abstractmethod def astype(self, dtype: DTypes, order: str = ..., casting: str = ..., subok: bool = ..., copy: bool = ...) -> "_Array4d": ... # These is actually a bit too strict: It's legal to say 'array4d + array5d' # That's kind of bad code though; it's better to write array5d + array4d. # We could relax this, but let's try the strict version. @abstractmethod def __add__(self: SelfT, other: Union[float, int, Array1d, Array2d, Array3d, "Array4d"]) -> SelfT: ... @abstractmethod def __sub__(self: SelfT, other: Union[float, int, Array1d, Array2d, Array3d, "Array4d"]) -> SelfT: ... @abstractmethod def __mul__(self: SelfT, other: Union[float, int, Array1d, Array2d, Array3d, "Array4d"]) -> SelfT: ... @abstractmethod def __pow__(self: SelfT, other: Union[float, int, Array1d, Array2d, Array3d, "Array4d"]) -> SelfT: ... @abstractmethod def __matmul__(self: SelfT, other: Union[float, int, Array1d, Array2d, Array3d, "Array4d"]) -> SelfT: ... @abstractmethod # These are not too strict though: you can't do += with higher dimensional. @abstractmethod def __iadd__(self, other: Union[float, int, Array1d, Array2d, Array3d, "Array4d"]): ... @abstractmethod def __isub__(self, other: Union[float, int, Array1d, Array2d, Array3d, "Array4d"]): ... @abstractmethod def __imul__(self, other: Union[float, int, Array1d, Array2d, Array3d, "Array4d"]): ... @abstractmethod def __ipow__(self, other: Union[float, int, Array1d, Array2d, Array3d, "Array4d"]): ... class Floats4d(_Array4d, _Floats): """4-dimensional array of floats.""" T: "Floats4d" @classmethod def __get_validators__(cls): """Runtime validation for pydantic.""" yield lambda v: validate_array(v, ndim=4, dtype="f") @abstractmethod def __iter__(self) -> Iterator[Floats3d]: ... @overload @abstractmethod def __getitem__(self, key: _4_KeyScalar) -> float: ... @overload @abstractmethod def __getitem__(self, key: _4_Key1d) -> Floats1d: ... @overload @abstractmethod def __getitem__(self, key: _4_Key2d) -> Floats2d: ... @overload @abstractmethod def __getitem__(self, key: _4_Key3d) -> Floats3d: ... @overload @abstractmethod def __getitem__(self, key: _4_Key4d) -> "Floats4d": ... @abstractmethod def __getitem__(self, key: _4_AllKeys) -> _F4_AllReturns: ... @overload @abstractmethod def __setitem__(self, key: _4_KeyScalar, value: float) -> None: ... @overload @abstractmethod def __setitem__(self, key: _4_Key1d, value: Floats1d) -> None: ... @overload @abstractmethod def __setitem__(self, key: _4_Key2d, value: Floats2d) -> None: ... @overload @abstractmethod def __setitem__(self, key: _4_Key3d, value: Floats3d) -> None: ... @overload @abstractmethod def __setitem__(self, key: _4_Key4d, value: "Floats4d") -> None: ... @abstractmethod def __setitem__(self, key: _4_AllKeys, value: _F4_AllReturns) -> None: ... @overload @abstractmethod def sum(self, *, keepdims: Tru, axis: _4_AllAx = None, out: Optional["Floats4d"] = None) -> "Floats4d": ... @overload @abstractmethod def sum(self, *, keepdims: Fal = False, axis: OneAx, out: Optional[Floats3d] = None) -> Floats3d: ... @overload @abstractmethod def sum(self, *, keepdims: Fal = False, axis: TwoAx, out: Optional[Floats2d] = None) -> Floats2d: ... @overload @abstractmethod def sum(self, *, keepdims: Fal = False, axis: ThreeAx, out: Optional[Floats1d] = None) -> Floats1d: ... @overload @abstractmethod def sum(self, *, keepdims: Fal = False, axis: Optional[FourAx], out = None) -> float: ... @abstractmethod def sum(self, *, keepdims: bool = False, axis: _4_AllAx = None, out: Union[None, Floats1d, Floats2d, Floats3d, "Floats4d"] = None) -> _4F_ReduceResults: ... class Ints4d(_Array4d, _Ints): """4-dimensional array of ints.""" T: "Ints4d" @classmethod def __get_validators__(cls): """Runtime validation for pydantic.""" yield lambda v: validate_array(v, ndim=4, dtype="i") @abstractmethod def __iter__(self) -> Iterator[Ints3d]: ... @overload @abstractmethod def __getitem__(self, key: _4_KeyScalar) -> int: ... @overload @abstractmethod def __getitem__(self, key: _4_Key1d) -> Ints1d: ... @overload @abstractmethod def __getitem__(self, key: _4_Key2d) -> Ints2d: ... @overload @abstractmethod def __getitem__(self, key: _4_Key3d) -> Ints3d: ... @overload @abstractmethod def __getitem__(self, key: _4_Key4d) -> "Ints4d": ... @abstractmethod def __getitem__(self, key: _4_AllKeys) -> _I4_AllReturns: ... @overload @abstractmethod def __setitem__(self, key: _4_KeyScalar, value: int) -> None: ... @overload @abstractmethod def __setitem__(self, key: _4_Key1d, value: Ints1d) -> None: ... @overload @abstractmethod def __setitem__(self, key: _4_Key2d, value: Ints2d) -> None: ... @overload @abstractmethod def __setitem__(self, key: _4_Key3d, value: Ints3d) -> None: ... @overload @abstractmethod def __setitem__(self, key: _4_Key4d, value: "Ints4d") -> None: ... @abstractmethod def __setitem__(self, key: _4_AllKeys, value: _I4_AllReturns) -> None: ... @overload @abstractmethod def sum(self, *, keepdims: Tru, axis: _4_AllAx = None, out: Optional["Ints4d"] = None) -> "Ints4d": ... @overload @abstractmethod def sum(self, *, keepdims: Fal = False, axis: OneAx, out: Optional[Ints3d] = None) -> Ints3d: ... @overload @abstractmethod def sum(self, *, keepdims: Fal = False, axis: TwoAx, out: Optional[Ints2d] = None) -> Ints2d: ... @overload @abstractmethod def sum(self, *, keepdims: Fal = False, axis: ThreeAx, out: Optional[Ints1d] = None) -> Ints1d: ... @overload @abstractmethod def sum(self, *, keepdims: Fal = False, axis: Optional[FourAx] = None, out = None) -> int: ... @abstractmethod def sum(self, *, keepdims: bool = False, axis: _4_AllAx = None, out: Optional[Union[Ints1d, Ints2d, Ints3d, "Ints4d"]] = None) -> _4I_ReduceResults: ... _DIn = TypeVar("_DIn") class Decorator(Protocol): """Protocol to mark a function as returning its child with identical signature.""" def __call__(self, name: str) -> Callable[[_DIn], _DIn]: ... # fmt: on class Generator(Iterator): """Custom generator type. Used to annotate function arguments that accept generators so they can be validated by pydantic (which doesn't support iterators/iterables otherwise). """ @classmethod def __get_validators__(cls): yield cls.validate @classmethod def validate(cls, v): if not hasattr(v, "__iter__") and not hasattr(v, "__next__"): raise TypeError("not a valid iterator") return v @dataclass class SizedGenerator: """A generator that has a __len__ and can repeatedly call the generator function. """ get_items: Callable[[], Generator] length: int def __len__(self): return self.length def __iter__(self): yield from self.get_items() @dataclass class Padded: """A batch of padded sequences, sorted by decreasing length. The data array is of shape (step, batch, ...). The auxiliary array size_at_t indicates the length of the batch at each timestep, so you can do data[:, :size_at_t[t]] to shrink the batch. The lengths array indicates the length of each row b, and the indices indicates the original ordering. """ data: Array3d size_at_t: Ints1d lengths: Ints1d indices: Ints1d def copy(self): return Padded( self.data.copy(), self.size_at_t.copy(), self.lengths.copy(), self.indices.copy(), ) def __len__(self) -> int: return self.lengths.shape[0] def __getitem__(self, index: Union[int, slice, Ints1d]) -> "Padded": if isinstance(index, int): # Slice to keep the dimensionality return Padded( self.data[:, index : index + 1], self.lengths[index : index + 1], self.lengths[index : index + 1], self.indices[index : index + 1], ) elif isinstance(index, slice): return Padded( self.data[:, index], self.lengths[index], self.lengths[index], self.indices[index], ) else: # If we get a sequence of indices, we need to be careful that # we maintain the length-sorting, while also keeping the mapping # back to the original order correct. sorted_index = list(sorted(index)) return Padded( self.data[sorted_index], self.size_at_t[sorted_index], self.lengths[sorted_index], self.indices[index], # Use original, to maintain order. ) @dataclass class Ragged: """A batch of concatenated sequences, that vary in the size of their first dimension. Ragged allows variable-length sequence data to be contiguous in memory, without padding. Indexing into Ragged is just like indexing into the *lengths* array, except it returns a Ragged object with the accompanying sequence data. For instance, you can write ragged[1:4] to get a Ragged object with sequences 1, 2 and 3. """ data: Array2d lengths: Ints1d data_shape: Tuple[int, ...] starts_ends: Optional[Ints1d] = None def __init__(self, data: _Array, lengths: Ints1d): self.lengths = lengths # Frustratingly, the -1 dimension doesn't work with 0 size... if data.size: self.data = cast(Array2d, data.reshape((data.shape[0], -1))) else: self.data = cast(Array2d, data.reshape((0, 0))) self.data_shape = (-1,) + data.shape[1:] @property def dataXd(self) -> ArrayXd: if self.data.size: reshaped = self.data.reshape(self.data_shape) else: reshaped = self.data.reshape((self.data.shape[0],) + self.data_shape[1:]) return cast(ArrayXd, reshaped) def __len__(self) -> int: return self.lengths.shape[0] def __getitem__(self, index: Union[int, slice, Array1d]) -> "Ragged": if isinstance(index, tuple): raise IndexError("Ragged arrays do not support 2d indexing.") starts = self._get_starts() ends = self._get_ends() if isinstance(index, int): s = starts[index] e = ends[index] return Ragged(self.data[s:e], self.lengths[index : index + 1]) elif isinstance(index, slice): lengths = self.lengths[index] if len(lengths) == 0: return Ragged(self.data[0:0].reshape(self.data_shape), lengths) start = starts[index][0] if index.start >= 1 else 0 end = ends[index][-1] return Ragged(self.data[start:end].reshape(self.data_shape), lengths) else: # There must be a way to do this "properly" :(. Sigh, hate numpy. xp = get_array_module(self.data) data = xp.vstack([self[int(i)].data for i in index]) return Ragged(data.reshape(self.data_shape), self.lengths[index]) def _get_starts_ends(self) -> Ints1d: if self.starts_ends is None: xp = get_array_module(self.lengths) self.starts_ends = xp.empty(self.lengths.size + 1, dtype="i") self.starts_ends[0] = 0 self.lengths.cumsum(out=self.starts_ends[1:]) return self.starts_ends def _get_starts(self) -> Ints1d: return self._get_starts_ends()[:-1] def _get_ends(self) -> Ints1d: return self._get_starts_ends()[1:] _P = TypeVar("_P", bound=Sequence) @dataclass class Pairs(Generic[_P]): """Dataclass for pairs of sequences that allows indexing into the sequences while keeping them aligned. """ one: _P two: _P def __getitem__(self, index) -> "Pairs[_P]": return Pairs(self.one[index], self.two[index]) def __len__(self) -> int: return len(self.one) @dataclass class ArgsKwargs: """A tuple of (args, kwargs) that can be spread into some function f: f(*args, **kwargs) """ args: Tuple[Any, ...] kwargs: Dict[str, Any] @classmethod def from_items(cls, items: Sequence[Tuple[Union[int, str], Any]]) -> "ArgsKwargs": """Create an ArgsKwargs object from a sequence of (key, value) tuples, such as produced by argskwargs.items(). Each key should be either a string or an integer. Items with int keys are added to the args list, and items with string keys are added to the kwargs list. The args list is determined by sequence order, not the value of the integer. """ args = [] kwargs = {} for key, value in items: if isinstance(key, int): args.append(value) else: kwargs[key] = value return cls(args=tuple(args), kwargs=kwargs) def keys(self) -> Iterable[Union[int, str]]: """Yield indices from self.args, followed by keys from self.kwargs.""" yield from range(len(self.args)) yield from self.kwargs.keys() def values(self) -> Iterable[Any]: """Yield elements of from self.args, followed by values from self.kwargs.""" yield from self.args yield from self.kwargs.values() def items(self) -> Iterable[Tuple[Union[int, str], Any]]: """Yield enumerate(self.args), followed by self.kwargs.items()""" yield from enumerate(self.args) yield from self.kwargs.items() @dataclass class Unserializable: """Wrap a value to prevent it from being serialized by msgpack.""" obj: Any def validate_array(obj, ndim=None, dtype=None): """Runtime validator for pydantic to validate array types.""" xp = get_array_module(obj) if not isinstance(obj, xp.ndarray): raise TypeError("not a valid numpy or cupy array") errors = [] if ndim is not None and obj.ndim != ndim: errors.append(f"wrong array dimensions (expected {ndim}, got {obj.ndim})") if dtype is not None: dtype_mapping = {"f": ["float32"], "i": ["int32", "int64", "uint32", "uint64"]} expected_types = dtype_mapping.get(dtype, []) if obj.dtype not in expected_types: expected = "/".join(expected_types) err = f"wrong array data type (expected {expected}, got {obj.dtype})" errors.append(err) if errors: raise ValueError(", ".join(errors)) return obj