rocm_jax/jax/numpy/__init__.pyi

from __future__ import annotations


import builtins
from typing import Any, Callable, Literal, NamedTuple, Optional, Sequence, TypeVar, Union, overload

from jax._src import core as _core
from jax._src import dtypes as _dtypes
from jax._src.lax.lax import PrecisionLike
from jax._src.lax.slicing import GatherScatterMode
from jax._src.numpy.index_tricks import _Mgrid, _Ogrid, CClass as _CClass, RClass as _RClass
from jax._src.typing import (
    Array, ArrayLike, DType, DTypeLike,
    DimSize, DuckTypedArray, Shape, DeprecatedArg
)
from jax.numpy import fft as fft, linalg as linalg
from jax.sharding import Sharding as _Sharding
import numpy as _np

_T = TypeVar('_T')

_Axis = Union[None, int, Sequence[int]]

# TODO(jakevdp): use xla_client.Device here
_Device = Any

ComplexWarning: type

_deprecations: dict[str, tuple[str, Any]]
def abs(x: ArrayLike, /) -> Array: ...
def absolute(x: ArrayLike, /) -> Array: ...
def acos(x: ArrayLike, /) -> Array: ...
def acosh(x: ArrayLike, /) -> Array: ...
def add(x: ArrayLike, y: ArrayLike, /) -> Array: ...
def amax(a: ArrayLike, axis: _Axis = ..., out: None = ...,
        keepdims: builtins.bool = ..., initial: Optional[ArrayLike] = ...,
        where: Optional[ArrayLike] = ...) -> Array: ...
def amin(a: ArrayLike, axis: _Axis = ..., out: None = ...,
        keepdims: builtins.bool = ..., initial: Optional[ArrayLike] = ...,
        where: Optional[ArrayLike] = ...) -> Array: ...
def all(a: ArrayLike, axis: _Axis = ..., out: None = ...,
        keepdims: builtins.bool = ..., *, where: Optional[ArrayLike] = ...) -> Array: ...
def allclose(a: ArrayLike, b: ArrayLike, rtol: ArrayLike = ...,
             atol: ArrayLike = ..., equal_nan: builtins.bool = ...) -> Array: ...
alltrue = all
def angle(z: ArrayLike, deg: builtins.bool = ...) -> Array: ...
def any(a: ArrayLike, axis: _Axis = ..., out: None = ...,
        keepdims: builtins.bool = ..., *, where: Optional[ArrayLike] = ...) -> Array: ...
def append(
    arr: ArrayLike, values: ArrayLike, axis: Optional[int] = ...
) -> Array: ...
def apply_along_axis(func1d: Callable, axis: int, arr: ArrayLike, *args,
                     **kwargs) -> Array: ...
def apply_over_axes(
    func: Callable, a: ArrayLike, axes: Sequence[int]
) -> Array: ...
def arange(
    start: DimSize,
    stop: Optional[DimSize] = ...,
    step: Optional[DimSize] = ...,
    dtype: Optional[DTypeLike] = ...,
) -> Array: ...
def arccos(x: ArrayLike, /) -> Array: ...
def arccosh(x: ArrayLike, /) -> Array: ...
def arcsin(x: ArrayLike, /) -> Array: ...
def arcsinh(x: ArrayLike, /) -> Array: ...
def arctan(x: ArrayLike, /) -> Array: ...
def arctan2(x: ArrayLike, y: ArrayLike, /) -> Array: ...
def arctanh(x: ArrayLike, /) -> Array: ...
def argmax(
    a: ArrayLike,
    axis: Optional[int] = ...,
    out: None = ...,
    keepdims: Optional[builtins.bool] = ...,
) -> Array: ...
def argmin(
    a: ArrayLike,
    axis: Optional[int] = ...,
    out: None = ...,
    keepdims: Optional[builtins.bool] = ...,
) -> Array: ...
def argpartition(a: ArrayLike, kth: int, axis: int = ...) -> Array: ...
def argsort(
    a: ArrayLike,
    axis: Optional[int] = ...,
    kind: str | None = ...,
    order: None = ...,
    *,
    stable: builtins.bool = ...,
    descending: builtins.bool = ...,
) -> Array: ...
def argwhere(
    a: ArrayLike,
    *,
    size: Optional[int] = ...,
    fill_value: Optional[ArrayLike] = ...,
) -> Array: ...
around = round
def array(object: Any, dtype: DTypeLike | None = ..., copy: builtins.bool = True,
          order: str | None = ..., ndmin: int = ...) -> Array: ...
def array_equal(
    a1: ArrayLike, a2: ArrayLike, equal_nan: builtins.bool = ...
) -> Array: ...
def array_equiv(a1: ArrayLike, a2: ArrayLike) -> Array: ...
array_repr = _np.array_repr
def array_split(
    ary: ArrayLike,
    indices_or_sections: Union[int, Sequence[int], ArrayLike],
    axis: int = ...,
) -> list[Array]: ...
array_str = _np.array_str
def asarray(
    a: Any, dtype: Optional[DTypeLike] = ..., order: Optional[str] = ...,
    *, copy: Optional[builtins.bool] = ...
) -> Array: ...
def asin(x: ArrayLike, /) -> Array: ...
def asinh(x: ArrayLike, /) -> Array: ...
def astype(a: ArrayLike, dtype: Optional[DTypeLike], /, *, copy: builtins.bool = ...) -> Array: ...
def atan(x: ArrayLike, /) -> Array: ...
def atan2(x: ArrayLike, y: ArrayLike, /) -> Array: ...
def atanh(x: ArrayLike, /) -> Array: ...
@overload
def atleast_1d() -> list[Array]: ...
@overload
def atleast_1d(x: ArrayLike, /) -> Array: ...
@overload
def atleast_1d(x: ArrayLike, y: ArrayLike, /, *arys: ArrayLike) -> list[Array]: ...

@overload
def atleast_2d() -> list[Array]: ...
@overload
def atleast_2d(x: ArrayLike, /) -> Array: ...
@overload
def atleast_2d(x: ArrayLike, y: ArrayLike, /, *arys: ArrayLike) -> list[Array]: ...

@overload
def atleast_3d() -> list[Array]: ...
@overload
def atleast_3d(x: ArrayLike, /) -> Array: ...
@overload
def atleast_3d(x: ArrayLike, y: ArrayLike, /, *arys: ArrayLike) -> list[Array]: ...

@overload
def average(a: ArrayLike, axis: _Axis = ..., weights: Optional[ArrayLike] = ...,
            returned: Literal[False] = False, keepdims: builtins.bool = False) -> Array: ...
@overload
def average(a: ArrayLike, axis: _Axis = ..., weights: Optional[ArrayLike] = ..., *,
            returned: Literal[True], keepdims: builtins.bool = False) -> tuple[Array, Array]: ...
@overload
def average(a: ArrayLike, axis: _Axis = ..., weights: Optional[ArrayLike] = ...,
            returned: builtins.bool = False, keepdims: builtins.bool = False) -> Union[Array, tuple[Array, Array]]: ...

def bartlett(M: int) -> Array: ...
bfloat16: Any
def bincount(x: ArrayLike, weights: Optional[ArrayLike] = ...,
             minlength: int = ..., *, length: Optional[int] = ...) -> Array: ...
def bitwise_and(x: ArrayLike, y: ArrayLike, /) -> Array: ...
def bitwise_count(x: ArrayLike, /) -> Array: ...
def bitwise_invert(x: ArrayLike, /) -> Array: ...
def bitwise_left_shift(x: ArrayLike, y: ArrayLike, /) -> Array: ...
def bitwise_not(x: ArrayLike, /) -> Array: ...
def bitwise_or(x: ArrayLike, y: ArrayLike, /) -> Array: ...
def bitwise_right_shift(x: ArrayLike, y: ArrayLike, /) -> Array: ...
def bitwise_xor(x: ArrayLike, y: ArrayLike, /) -> Array: ...
def blackman(M: int) -> Array: ...
def block(arrays: Union[ArrayLike, Sequence[ArrayLike], Sequence[Sequence[ArrayLike]]]) -> Array: ...
bool: Any
bool_: Any
def broadcast_arrays(*args: ArrayLike) -> list[Array]: ...

@overload
def broadcast_shapes(*shapes: Sequence[int]) -> tuple[int, ...]: ...

@overload
def broadcast_shapes(*shapes: Sequence[Union[int, _core.Tracer]]
                     ) -> tuple[Union[int, _core.Tracer], ...]: ...

def broadcast_to(array: ArrayLike, shape: DimSize | Shape) -> Array: ...
c_: _CClass
can_cast = _np.can_cast
def cbrt(x: ArrayLike, /) -> Array: ...
cdouble: Any
def ceil(x: ArrayLike, /) -> Array: ...
character = _np.character
def choose(a: ArrayLike, choices: Sequence[ArrayLike],
           out: None = ..., mode: str = ...) -> Array: ...
def clip(
    x: ArrayLike | None = ...,
    /,
    min: Optional[ArrayLike] = ...,
    max: Optional[ArrayLike] = ...,
    a: ArrayLike | DeprecatedArg | None = ...,
    a_min: ArrayLike | DeprecatedArg | None = ...,
    a_max: ArrayLike | DeprecatedArg | None = ...
) -> Array: ...
def column_stack(
    tup: Union[_np.ndarray, Array, Sequence[ArrayLike]]
) -> Array: ...
complex128: Any
complex64: Any
complex_: Any
complexfloating = _np.complexfloating
def compress(condition: ArrayLike, a: ArrayLike, axis: Optional[int] = ...,
             out: None = ...) -> Array: ...
def concat(arrays: Sequence[ArrayLike], /, *, axis: int | None = 0) -> Array: ...
def concatenate(
    arrays: Union[_np.ndarray, Array, Sequence[ArrayLike]],
    axis: Optional[int] = ...,
    dtype: Optional[DTypeLike] = ...,
) -> Array: ...
def conjugate(x: ArrayLike, /) -> Array: ...
conj = conjugate
def convolve(a: ArrayLike, v: ArrayLike, mode: str = ..., *,
             precision: PrecisionLike = ...,
             preferred_element_type: Optional[DTypeLike] = ...) -> Array: ...
def copy(a: ArrayLike, order: Optional[str] = ...) -> Array: ...
def copysign(x: ArrayLike, y: ArrayLike, /) -> Array: ...
def corrcoef(x: ArrayLike, y: Optional[ArrayLike] = ..., rowvar: builtins.bool = ...) -> Array: ...
def correlate(a: ArrayLike, v: ArrayLike, mode: str = ..., *,
              precision: PrecisionLike = ...,
              preferred_element_type: Optional[DTypeLike] = ...) -> Array: ...
def cos(x: ArrayLike, /) -> Array: ...
def cosh(x: ArrayLike, /) -> Array: ...
def count_nonzero(a: ArrayLike, axis: _Axis = ...,
                  keepdims: builtins.bool = ...) -> Array: ...
def cov(m: ArrayLike, y: Optional[ArrayLike] = ..., rowvar: builtins.bool = ...,
        bias: builtins.bool = ..., ddof: Optional[int] = ...,
        fweights: Optional[ArrayLike] = ...,
        aweights: Optional[ArrayLike] = ...) -> Array: ...
def cross(
    a: ArrayLike,
    b: ArrayLike,
    axisa: int = -1,
    axisb: int = -1,
    axisc: int = -1,
    axis: Optional[int] = ...,
) -> Array: ...
csingle: Any
def cumprod(a: ArrayLike, axis: _Axis = ..., dtype: DTypeLike = ...,
            out: None = ...) -> Array: ...
cumproduct = cumprod
def cumsum(a: ArrayLike, axis: _Axis = ..., dtype: DTypeLike = ...,
           out: None = ...) -> Array: ...

def deg2rad(x: ArrayLike, /) -> Array: ...
degrees = rad2deg
def delete(
    arr: ArrayLike,
    obj: Union[ArrayLike, slice],
    axis: Optional[int] = ...,
    *,
    assume_unique_indices: builtins.bool = ...,
) -> Array: ...
def diag(v: ArrayLike, k: int = 0) -> Array: ...
def diag_indices(n: int, ndim: int = ...) -> tuple[Array, ...]: ...
def diag_indices_from(arr: ArrayLike) -> tuple[Array, ...]: ...
def diagflat(v: ArrayLike, k: int = 0) -> Array: ...
def diagonal(
    a: ArrayLike, offset: ArrayLike = ..., axis1: int = ..., axis2: int = ...
): ...
def diff(a: ArrayLike, n: int = ..., axis: int = ...,
         prepend: Optional[ArrayLike] = ...,
         append: Optional[ArrayLike] = ...) -> Array: ...
def digitize(x: ArrayLike, bins: ArrayLike, right: builtins.bool = ...) -> Array: ...
divide = true_divide
def divmod(x: ArrayLike, y: ArrayLike, /) -> tuple[Array, Array]: ...
def dot(
    a: ArrayLike, b: ArrayLike, *, precision: PrecisionLike = ...,
    preferred_element_type: Optional[DTypeLike] = ...) -> Array: ...
double: Any
def dsplit(
    ary: ArrayLike, indices_or_sections: Union[int, ArrayLike]
) -> list[Array]: ...

def dstack(tup: Union[_np.ndarray, Array, Sequence[ArrayLike]],
           dtype: Optional[DTypeLike] = ...) -> Array: ...
dtype = _np.dtype
e: float
def ediff1d(ary: ArrayLike, to_end: Optional[ArrayLike] = ...,
            to_begin: Optional[ArrayLike] = ...) -> Array: ...
@overload
def einsum(
    subscript: str, /,
    *operands: ArrayLike,
    out: None = ...,
    optimize: str = "optimal",
    precision: PrecisionLike = ...,
    preferred_element_type: Optional[DTypeLike] = ...,
    _use_xeinsum: builtins.bool = False,
    _dot_general: Callable[..., Array] = ...,
) -> Array: ...

@overload
def einsum(
    arr: ArrayLike,
    axes: Sequence[Any], /,
    *operands: Union[ArrayLike, Sequence[Any]],
    out: None = ...,
    optimize: str = "optimal",
    precision: PrecisionLike = ...,
    preferred_element_type: Optional[DTypeLike] = ...,
    _use_xeinsum: builtins.bool = False,
    _dot_general: Callable[..., Array] = ...,
) -> Array: ...
@overload
def einsum(
    subscripts, /,
    *operands,
    out: None = ...,
    optimize: str = ...,
    precision: PrecisionLike = ...,
    preferred_element_type: Optional[DTypeLike] = ...,
    _use_xeinsum: builtins.bool = ...,
    _dot_general: Callable[..., Array] = ...,
) -> Array: ...

def einsum_path(subscripts, *operands, optimize =  ...): ...
def empty(shape: Any, dtype: Optional[DTypeLike] = ...,
          device: Optional[Union[_Device, _Sharding]] = ...) -> Array: ...
def empty_like(prototype: Union[ArrayLike, DuckTypedArray],
               dtype: Optional[DTypeLike] = ...,
               shape: Any = ..., *,
               device: Optional[Union[_Device, _Sharding]] = ...) -> Array: ...
def equal(x: ArrayLike, y: ArrayLike, /) -> Array: ...
euler_gamma: float
def exp(x: ArrayLike, /) -> Array: ...
def exp2(x: ArrayLike, /) -> Array: ...
def expand_dims(a: ArrayLike, axis: Union[int, Sequence[int]]) -> Array: ...
def expm1(x: ArrayLike, /) -> Array: ...
def extract(condition: ArrayLike, arr: ArrayLike) -> Array: ...
def eye(N: DimSize, M: Optional[DimSize] = ..., k: int = ...,
        dtype: Optional[DTypeLike] = ...) -> Array: ...
def fabs(x: ArrayLike, /) -> Array: ...
finfo = _dtypes.finfo
def fix(x: ArrayLike, out: None = ...) -> Array: ...
def flatnonzero(
    a: ArrayLike,
    *,
    size: Optional[int] = ...,
    fill_value: Union[None, ArrayLike, tuple[ArrayLike]] = ...,
) -> Array: ...
flexible = _np.flexible
def flip(
    m: ArrayLike, axis: Optional[Union[int, Sequence[int]]] = ...
) -> Array: ...

def fliplr(m: ArrayLike) -> Array: ...
def flipud(m: ArrayLike) -> Array: ...
float16: Any
float32: Any
float64: Any
float8_e4m3b11fnuz: Any
float8_e4m3fn: Any
float8_e4m3fnuz: Any
float8_e5m2: Any
float8_e5m2fnuz: Any
float_: Any
def float_power(x: ArrayLike, y: ArrayLike, /) -> Array: ...
floating = _np.floating
def floor(x: ArrayLike, /) -> Array: ...
def floor_divide(x: ArrayLike, y: ArrayLike, /) -> Array: ...
def fmax(x: ArrayLike, y: ArrayLike, /) -> Array: ...
def fmin(x: ArrayLike, y: ArrayLike, /) -> Array: ...
def fmod(x: ArrayLike, y: ArrayLike, /) -> Array: ...
def frexp(x: ArrayLike, /) -> tuple[Array, Array]: ...
def from_dlpack(x: Any) -> Array: ...
def frombuffer(buffer: Union[bytes, Any], dtype: DTypeLike = ...,
               count: int = ..., offset: int = ...) -> Array: ...
def fromfile(*args, **kwargs): ...
def fromfunction(function: Callable[..., Array], shape: Any,
                 *, dtype: DTypeLike = ..., **kwargs) -> Array: ...
def fromiter(*args, **kwargs): ...
def fromstring(
    string: str, dtype: DTypeLike = ..., count: int = ..., *, sep: str
) -> Array: ...
def full(shape: Any, fill_value: ArrayLike,
         dtype: Optional[DTypeLike] = ..., *,
         device: Optional[Union[_Device, _Sharding]] = ...) -> Array: ...
def full_like(a: Union[ArrayLike, DuckTypedArray],
              fill_value: ArrayLike, dtype: Optional[DTypeLike] = ...,
              shape: Any = ..., *,
              device: Optional[Union[_Device, _Sharding]] = ...) -> Array: ...
def gcd(x1: ArrayLike, x2: ArrayLike) -> Array: ...
generic = _np.generic
def geomspace(
    start: ArrayLike,
    stop: ArrayLike,
    num: int = ...,
    endpoint: builtins.bool = ...,
    dtype: Optional[DTypeLike] = ...,
    axis: int = ...,
) -> Array: ...
get_printoptions = _np.get_printoptions
def gradient(f: ArrayLike, *varargs: ArrayLike,
             axis: Optional[Union[int, Sequence[int]]] = ...,
             edge_order: Optional[int] = ...) -> Union[Array, list[Array]]: ...
def greater(x: ArrayLike, y: ArrayLike, /) -> Array: ...
def greater_equal(x: ArrayLike, y: ArrayLike, /) -> Array: ...
def hamming(M: int) -> Array: ...
def hanning(M: int) -> Array: ...
def heaviside(x: ArrayLike, y: ArrayLike, /) -> Array: ...
def histogram(a: ArrayLike, bins: ArrayLike = ...,
              range: Optional[Sequence[ArrayLike]] = ...,
              weights: Optional[ArrayLike] = ...,
              density: Optional[builtins.bool] = ...) -> tuple[Array, Array]: ...
def histogram2d(
    x: ArrayLike,
    y: ArrayLike,
    bins: Union[ArrayLike, Sequence[ArrayLike]] = ...,
    range: Optional[Sequence[Union[None, Array, Sequence[ArrayLike]]]] = ...,
    weights: Optional[ArrayLike] = ...,
    density: Optional[builtins.bool] = ...,
) -> tuple[Array, Array, Array]: ...
def histogram_bin_edges(a: ArrayLike, bins: ArrayLike = ...,
                        range: Union[None, Array, Sequence[ArrayLike]] = ...,
                        weights: Optional[ArrayLike] = ...) -> Array: ...
def histogramdd(
    sample: ArrayLike,
    bins: Union[ArrayLike, Sequence[ArrayLike]] = ...,
    range: Optional[Sequence[Union[None, Array, Sequence[ArrayLike]]]] = ...,
    weights: Optional[ArrayLike] = ...,
    density: Optional[builtins.bool] = ...,
) -> tuple[Array, list[Array]]: ...
def hsplit(
    ary: ArrayLike, indices_or_sections: Union[int, ArrayLike]
) -> list[Array]: ...
def hstack(tup: Union[_np.ndarray, Array, Sequence[ArrayLike]],
           dtype: Optional[DTypeLike] = ...) -> Array: ...
def hypot(x: ArrayLike, y: ArrayLike, /) -> Array: ...
def i0(x: ArrayLike) -> Array: ...
def identity(n: DimSize, dtype: Optional[DTypeLike] = ...) -> Array: ...
iinfo = _dtypes.iinfo
def imag(x: ArrayLike, /) -> Array: ...
index_exp = _np.index_exp

@overload
def indices(dimensions: Sequence[int], dtype: DTypeLike = int32,
            sparse: Literal[False] = False) -> Array: ...
@overload
def indices(dimensions: Sequence[int], dtype: DTypeLike = int32,
            *, sparse: Literal[True]) -> tuple[Array, ...]: ...
@overload
def indices(dimensions: Sequence[int], dtype: DTypeLike = int32,
            sparse: builtins.bool = False) -> Union[Array, tuple[Array, ...]]: ...

inexact = _np.inexact
inf: float
def inner(
    a: ArrayLike, b: ArrayLike, *, precision: PrecisionLike = ...,
    preferred_element_type: Optional[DTypeLike] = ...) -> Array: ...
def insert(arr: ArrayLike, obj: Union[ArrayLike, slice], values: ArrayLike,
           axis: Optional[int] = ...) -> Array: ...
int16: Any
int32: Any
int4: Any
int64: Any
int8: Any
int_: Any
integer = _np.integer
def interp(x: ArrayLike, xp: ArrayLike, fp: ArrayLike,
           left: Union[ArrayLike, str, None] = ...,
           right: Union[ArrayLike, str, None] = ...,
           period: Optional[ArrayLike] = ...) -> Array: ...
def intersect1d(ar1: ArrayLike, ar2: ArrayLike, assume_unique: builtins.bool = ...,
                return_indices: builtins.bool = ...) -> Union[Array, tuple[Array, Array, Array]]: ...
def invert(x: ArrayLike, /) -> Array: ...
def isclose(a: ArrayLike, b: ArrayLike, rtol: ArrayLike = ...,
            atol: ArrayLike = ..., equal_nan: builtins.bool = ...) -> Array: ...
def iscomplex(m: ArrayLike) -> Array: ...
def iscomplexobj(x: Any) -> builtins.bool: ...
def isdtype(dtype: DTypeLike, kind: Union[DType, str, tuple[Union[DType, str], ...]]) -> builtins.bool: ...
def isfinite(x: ArrayLike, /) -> Array: ...
def isin(element: ArrayLike, test_elements: ArrayLike,
         assume_unique: builtins.bool = ..., invert: builtins.bool = ...) -> Array: ...
def isinf(x: ArrayLike, /) -> Array: ...
def isnan(x: ArrayLike, /) -> Array: ...
def isneginf(x: ArrayLike, /) -> Array: ...
def isposinf(x: ArrayLike, /) -> Array: ...
def isreal(m: ArrayLike) -> Array: ...
def isrealobj(x: Any) -> builtins.bool: ...
def isscalar(element: Any) -> builtins.bool: ...
def issubdtype(arg1: DTypeLike, arg2: DTypeLike) -> builtins.bool: ...
iterable = _np.iterable
def ix_(*args: ArrayLike) -> tuple[Array, ...]: ...
def kaiser(M: int, beta: ArrayLike) -> Array: ...
def kron(a: ArrayLike, b: ArrayLike) -> Array: ...
def lcm(x1: ArrayLike, x2: ArrayLike) -> Array: ...
def ldexp(x: ArrayLike, y: ArrayLike, /) -> Array: ...
def left_shift(x: ArrayLike, y: ArrayLike, /) -> Array: ...
def less(x: ArrayLike, y: ArrayLike, /) -> Array: ...
def less_equal(x: ArrayLike, y: ArrayLike, /) -> Array: ...
def lexsort(keys: Sequence[ArrayLike], axis: int = ...) -> Array: ...

@overload
def linspace(start: ArrayLike, stop: ArrayLike, num: int = 50,
             endpoint: builtins.bool = True, retstep: Literal[False] = False,
             dtype: Optional[DTypeLike] = ...,
             axis: int = 0) -> Array: ...
@overload
def linspace(start: ArrayLike, stop: ArrayLike, num: int,
             endpoint: builtins.bool, retstep: Literal[True],
             dtype: Optional[DTypeLike] = ...,
             axis: int = 0) -> tuple[Array, Array]: ...
@overload
def linspace(start: ArrayLike, stop: ArrayLike, num: int = 50,
             endpoint: builtins.bool = True, *, retstep: Literal[True],
             dtype: Optional[DTypeLike] = ...,
             axis: int = 0) -> tuple[Array, Array]: ...
@overload
def linspace(start: ArrayLike, stop: ArrayLike, num: int = 50,
             endpoint: builtins.bool = True, retstep: builtins.bool = False,
             dtype: Optional[DTypeLike] = ...,
             axis: int = 0) -> Union[Array, tuple[Array, Array]]: ...

def load(*args: Any, **kwargs: Any) -> Array: ...
def log(x: ArrayLike, /) -> Array: ...
def log10(x: ArrayLike, /) -> Array: ...
def log1p(x: ArrayLike, /) -> Array: ...
def log2(x: ArrayLike, /) -> Array: ...
def logaddexp(x: ArrayLike, y: ArrayLike, /) -> Array: ...
def logaddexp2(x: ArrayLike, y: ArrayLike, /) -> Array: ...
def logical_and(x: ArrayLike, y: ArrayLike, /) -> Array: ...
def logical_not(x: ArrayLike, /) -> Array: ...
def logical_or(x: ArrayLike, y: ArrayLike, /) -> Array: ...
def logical_xor(x: ArrayLike, y: ArrayLike, /) -> Array: ...
def logspace(start: ArrayLike, stop: ArrayLike, num: int = ...,
             endpoint: builtins.bool = ..., base: ArrayLike = ...,
             dtype: Optional[DTypeLike] = ..., axis: int = ...) -> Array: ...
def mask_indices(
    n: int, mask_func: Callable, k: int = ...
) -> tuple[Array, ...]: ...
def matmul(
    a: ArrayLike, b: ArrayLike, *, precision: PrecisionLike = ...,
    preferred_element_type: Optional[DTypeLike] = ...) -> Array: ...
def matrix_transpose(x: ArrayLike, /) -> Array: ...
max = amax
def maximum(x: ArrayLike, y: ArrayLike, /) -> Array: ...
def mean(a: ArrayLike, axis: _Axis = ..., dtype: DTypeLike = ...,
         out: None = ..., keepdims: builtins.bool = ..., *,
         where: Optional[ArrayLike] = ...) -> Array: ...
def median(a: ArrayLike, axis: Optional[Union[int, tuple[int, ...]]] = ...,
           out: None = ..., overwrite_input: builtins.bool = ...,
           keepdims: builtins.bool = ...) -> Array: ...
def meshgrid(*xi: ArrayLike, copy: builtins.bool = ..., sparse: builtins.bool = ...,
             indexing: str = ...) -> list[Array]: ...
mgrid: _Mgrid
min = amin
def minimum(x: ArrayLike, y: ArrayLike, /) -> Array: ...
def mod(x: ArrayLike, y: ArrayLike, /) -> Array: ...
def modf(x: ArrayLike, /, out=None) -> tuple[Array, Array]: ...
def moveaxis(a: ArrayLike, source: Union[int, Sequence[int]],
             destination: Union[int, Sequence[int]]) -> Array: ...
def multiply(x: ArrayLike, y: ArrayLike, /) -> Array: ...
nan: float
def nan_to_num(x: ArrayLike, copy: builtins.bool = ..., nan: ArrayLike = ...,
               posinf: Optional[ArrayLike] = ...,
               neginf: Optional[ArrayLike] = ...) -> Array: ...
def nanargmax(
    a: ArrayLike,
    axis: Optional[int] = ...,
    out: None = ...,
    keepdims: Optional[builtins.bool] = ...,
) -> Array: ...
def nanargmin(
    a: ArrayLike,
    axis: Optional[int] = ...,
    out: None = ...,
    keepdims: Optional[builtins.bool] = ...,
) -> Array: ...
def nancumprod(a: ArrayLike, axis: _Axis = ..., dtype: DTypeLike = ...,
               out: None = ...) -> Array: ...
def nancumsum(a: ArrayLike, axis: _Axis = ..., dtype: DTypeLike = ...,
               out: None = ...) -> Array: ...
def nanmax(a: ArrayLike, axis: _Axis = ..., out: None = ...,
           keepdims: builtins.bool = ..., initial: Optional[ArrayLike] = ...,
           where: Optional[ArrayLike] = ...) -> Array: ...
def nanmean(a: ArrayLike, axis: _Axis = ..., dtype: DTypeLike = ...,
            out: None = ...,
            keepdims: builtins.bool = ...,
            where: Optional[ArrayLike] = ...) -> Array: ...
def nanmedian(a: ArrayLike, axis: Optional[Union[int, tuple[int, ...]]] = ...,
              out: None = ..., overwrite_input: builtins.bool = ...,
              keepdims: builtins.bool = ...) -> Array: ...
def nanmin(a: ArrayLike, axis: _Axis = ..., out: None = ...,
           keepdims: builtins.bool = ..., initial: Optional[ArrayLike] = ...,
           where: Optional[ArrayLike] = ...) -> Array: ...
def nanpercentile(a: ArrayLike, q: ArrayLike,
                  axis: Optional[Union[int, tuple[int, ...]]] = ...,
                  out: None = ..., overwrite_input: builtins.bool = ..., method: str = ...,
                  keepdims: builtins.bool = ..., interpolation: None = ...) -> Array: ...
def nanprod(a: ArrayLike, axis: _Axis = ..., dtype: DTypeLike = ...,
            out: None = ...,
            keepdims: builtins.bool = ..., initial: Optional[ArrayLike] = ...,
            where: Optional[ArrayLike] = ...) -> Array: ...
def nanquantile(a: ArrayLike, q: ArrayLike, axis: Optional[Union[int, tuple[int, ...]]] = ...,
                out: None = ..., overwrite_input: builtins.bool = ..., method: str = ...,
                keepdims: builtins.bool = ..., interpolation: None = ...) -> Array: ...
def nanstd(a: ArrayLike, axis: _Axis = ..., dtype: DTypeLike = ..., out: None = ...,
           ddof: int = ..., keepdims: builtins.bool = ...,
           where: Optional[ArrayLike] = ...) -> Array: ...
def nansum(a: ArrayLike, axis: _Axis = ..., dtype: DTypeLike = ...,
           out: None = ..., keepdims: builtins.bool = ...,
           initial: Optional[ArrayLike] = ...,
           where: Optional[ArrayLike] = ...) -> Array: ...
def nanvar(a: ArrayLike, axis: _Axis = ..., dtype: DTypeLike = ...,
           out: None = ...,
           ddof: int = 0, keepdims: builtins.bool = False,
           where: Optional[ArrayLike] = ...) -> Array: ...
ndarray = Array
ndim = _np.ndim
def negative(x: ArrayLike, /) -> Array: ...
newaxis = ...
def nextafter(x: ArrayLike, y: ArrayLike, /) -> Array: ...
def nonzero(a: ArrayLike, *, size: Optional[int] = ...,
            fill_value: Union[None, ArrayLike, tuple[ArrayLike, ...]] = ...
    ) -> tuple[Array, ...]: ...
def not_equal(x: ArrayLike, y: ArrayLike, /) -> Array: ...
number = _np.number
object_ = _np.object_
ogrid: _Ogrid
def ones(shape: Any, dtype: Optional[DTypeLike] = ...,
         device: Optional[Union[_Device, _Sharding]] = ...) -> Array: ...
def ones_like(a: Union[ArrayLike, DuckTypedArray],
              dtype: Optional[DTypeLike] = ...,
              shape: Any = ..., *,
              device: Optional[Union[_Device, _Sharding]] = ...) -> Array: ...
def outer(a: ArrayLike, b: Array, out: None = ...) -> Array: ...
def packbits(
    a: ArrayLike, axis: Optional[int] = ..., bitorder: str = ...
) -> Array: ...

PadValueLike = Union[_T, Sequence[_T], Sequence[Sequence[_T]]]
def pad(array: ArrayLike, pad_width: PadValueLike[int | Array | _np.ndarray],
        mode: Union[str, Callable[..., Any]] = ..., **kwargs) -> Array: ...

def partition(a: ArrayLike, kth: int, axis: int = ...) -> Array: ...
def percentile(a: ArrayLike, q: ArrayLike,
               axis: Optional[Union[int, tuple[int, ...]]] = ...,
               out: None = ..., overwrite_input: builtins.bool = ..., method: str = ...,
               keepdims: builtins.bool = ..., interpolation: None = ...) -> Array: ...
def permute_dims(x: ArrayLike, /, axes: tuple[int, ...]) -> Array: ...
pi: float
def piecewise(x: ArrayLike, condlist: Union[Array, Sequence[ArrayLike]],
              funclist: Sequence[Union[ArrayLike, Callable[..., Array]]],
              *args, **kw) -> Array: ...
def place(arr: ArrayLike, mask: ArrayLike, vals: ArrayLike, *,
          inplace: builtins.bool = ...) -> Array: ...
def poly(seq_of_zeros: Array) -> Array: ...
def polyadd(a1: Array, a2: Array) -> Array: ...
def polyder(p: Array, m: int = ...) -> Array: ...
def polydiv(u: ArrayLike, v: ArrayLike, *, trim_leading_zeros: builtins.bool = ...) -> tuple[Array, Array]: ...
def polyfit(x: Array, y: Array, deg: int, rcond: Optional[float] = ...,
            full: builtins.bool = ..., w: Optional[Array] = ..., cov: builtins.bool = ...
            ) -> Union[Array, tuple[Array, ...]]: ...
def polyint(p: Array, m: int = ..., k: Optional[int] = ...) -> Array: ...
def polymul(a1: ArrayLike, a2: ArrayLike, *, trim_leading_zeros: builtins.bool = ...) -> Array: ...
def polysub(a1: Array, a2: Array) -> Array: ...
def polyval(p: Array, x: Array, *, unroll: int = ...) -> Array: ...
def positive(x: ArrayLike, /) -> Array: ...
def pow(x: ArrayLike, y: ArrayLike, /) -> Array: ...
def power(x: ArrayLike, y: ArrayLike, /) -> Array: ...
printoptions = _np.printoptions
def prod(a: ArrayLike, axis: _Axis = ..., dtype: DTypeLike = ...,
         out: None = ..., keepdims: builtins.bool = ...,
         initial: Optional[ArrayLike] = ..., where: Optional[ArrayLike] = ...,
         promote_integers: builtins.bool = ...) -> Array: ...
product = prod
promote_types = _np.promote_types
def ptp(a: ArrayLike, axis: _Axis = ..., out: None = ...,
        keepdims: builtins.bool = ...) -> Array: ...
def put(a: ArrayLike, ind: ArrayLike, v: ArrayLike,
        mode: str | None = ..., *, inplace: builtins.bool = ...) -> Array: ...
def quantile(a: ArrayLike, q: ArrayLike, axis: Optional[Union[int, tuple[int, ...]]] = ...,
             out: None = ..., overwrite_input: builtins.bool = ..., method: str = ...,
             keepdims: builtins.bool = ..., interpolation: None = ...) -> Array: ...
r_: _RClass
def rad2deg(x: ArrayLike, /) -> Array: ...
radians = deg2rad
def ravel(a: ArrayLike, order: str = ...) -> Array: ...
def ravel_multi_index(multi_index: Sequence[ArrayLike], dims: Sequence[int],
                      mode: str = ..., order: str = ...) -> Array: ...
def real(x: ArrayLike, /) -> Array: ...
def reciprocal(x: ArrayLike, /) -> Array: ...
register_jax_array_methods: Any
def remainder(x: ArrayLike, y: ArrayLike, /) -> Array: ...
def repeat(a: ArrayLike, repeats: ArrayLike, axis: Optional[int] = ..., *,
           total_repeat_length: Optional[int] = ...) -> Array: ...
def reshape(
    a: ArrayLike, newshape: Union[DimSize, Shape], order: str = ...
) -> Array: ...

def resize(a: ArrayLike, new_shape: Shape) -> Array: ...
def result_type(*args: Any) -> DType: ...
def right_shift(x: ArrayLike, y: ArrayLike, /) -> Array: ...
def rint(x: ArrayLike, /) -> Array: ...
def roll(a: ArrayLike, shift: Union[ArrayLike, Sequence[int]],
         axis: Optional[Union[int, Sequence[int]]] = ...) -> Array: ...
def rollaxis(a: ArrayLike, axis: int, start: int = 0) -> Array: ...
def roots(p: ArrayLike, *, strip_zeros: builtins.bool = ...) -> Array: ...
def rot90(m: ArrayLike, k: int = ..., axes: tuple[int, int] = ...) -> Array: ...
def round(a: ArrayLike, decimals: int = ..., out: None = ...) -> Array: ...
round_ = round
s_ = _np.s_
save = _np.save
savez = _np.savez
def searchsorted(a: ArrayLike, v: ArrayLike, side: str = ...,
                 sorter: None = ..., *, method: str = ...) -> Array: ...
def select(
    condlist: Sequence[ArrayLike],
    choicelist: Sequence[ArrayLike],
    default: ArrayLike = ...,
) -> Array: ...
set_printoptions = _np.set_printoptions
def setdiff1d(
    ar1: ArrayLike,
    ar2: ArrayLike,
    assume_unique: builtins.bool = ...,
    *,
    size: Optional[int] = ...,
    fill_value: Optional[ArrayLike] = ...,
) -> Array: ...
def setxor1d(ar1: ArrayLike, ar2: ArrayLike, assume_unique: builtins.bool = ...) -> Array: ...
shape = _np.shape
def sign(x: ArrayLike, /) -> Array: ...
def signbit(x: ArrayLike, /) -> Array: ...
signedinteger = _np.signedinteger
def sin(x: ArrayLike, /) -> Array: ...
def sinc(x: ArrayLike, /) -> Array: ...
single: Any
def sinh(x: ArrayLike, /) -> Array: ...
size = _np.size
sometrue = any
def sort(
    a: ArrayLike,
    axis: Optional[int] = ...,
    kind: str | None = ...,
    order: None = ...,
    *,
    stable: builtins.bool = ...,
    descending: builtins.bool = ...,
) -> Array: ...
def sort_complex(a: ArrayLike) -> Array: ...
def split(
    ary: ArrayLike,
    indices_or_sections: Union[int, Sequence[int], ArrayLike],
    axis: int = ...,
) -> list[Array]: ...

def sqrt(x: ArrayLike, /) -> Array: ...
def square(x: ArrayLike, /) -> Array: ...
def squeeze(
    a: ArrayLike, axis: Optional[Union[int, Sequence[int]]] = ...
) -> Array: ...
def stack(
    arrays: Union[_np.ndarray, Array, Sequence[ArrayLike]],
    axis: int = ...,
    out: None = ...,
    dtype: Optional[DTypeLike] = ...,
) -> Array: ...
def std(a: ArrayLike, axis: _Axis = ..., dtype: DTypeLike = ...,
        out: None = ..., ddof: int = ..., keepdims: builtins.bool = ..., *,
        where: Optional[ArrayLike] = ...) -> Array: ...
def subtract(x: ArrayLike, y: ArrayLike, /) -> Array: ...
def sum(
    a: ArrayLike,
    axis: _Axis = ...,
    dtype: DTypeLike = ...,
    out: None = ...,
    keepdims: builtins.bool = ...,
    initial: Optional[ArrayLike] = ...,
    where: Optional[ArrayLike] = ...,
    promote_integers: builtins.bool = ...,
) -> Array: ...
def swapaxes(a: ArrayLike, axis1: int, axis2: int) -> Array: ...
def take(
    a: ArrayLike,
    indices: ArrayLike,
    axis: Optional[int] = ...,
    out: None = ...,
    mode: Optional[str] = ...,
    unique_indices: builtins.bool = ...,
    indices_are_sorted: builtins.bool = ...,
    fill_value: Optional[ArrayLike] = ...,
) -> Array: ...
def take_along_axis(
    arr: ArrayLike,
    indices: ArrayLike,
    axis: Optional[int],
    mode: Optional[Union[str, GatherScatterMode]] = ...,
) -> Array: ...
def tan(x: ArrayLike, /) -> Array: ...
def tanh(x: ArrayLike, /) -> Array: ...
def tensordot(a: ArrayLike, b: ArrayLike,
              axes: Union[int, Sequence[int], Sequence[Sequence[int]]] = ...,
              *, precision: PrecisionLike = ...,
              preferred_element_type: Optional[DTypeLike] = ...) -> Array: ...
def tile(A: ArrayLike, reps: Union[DimSize, Sequence[DimSize]]) -> Array: ...
def trace(a: ArrayLike, offset: int = ..., axis1: int = ..., axis2: int = ...,
          dtype: Optional[DTypeLike] = ..., out: None = ...) -> Array: ...
def transpose(a: ArrayLike, axes: Optional[Sequence[int]] = ...) -> Array: ...
def trapezoid(y: ArrayLike, x: ArrayLike | None = None, dx: ArrayLike = ...,
              axis: int = ...) -> Array: ...
def tri(
    N: int, M: Optional[int] = ..., k: int = ..., dtype: DTypeLike = ...
) -> Array: ...
def tril(m: ArrayLike, k: int = ...) -> Array: ...
def tril_indices(
    n: int, k: int = ..., m: Optional[int] = ...
) -> tuple[Array, Array]: ...
def tril_indices_from(arr: ArrayLike, k: int = ...) -> tuple[Array, Array]: ...
def fill_diagonal(a: ArrayLike, val: ArrayLike, wrap: builtins.bool = ..., *, inplace: builtins.bool = ...) -> Array: ...
def trim_zeros(filt: ArrayLike, trim: str = ...) -> Array: ...
def triu(m: ArrayLike, k: int = ...) -> Array: ...
def triu_indices(
    n: int, k: int = ..., m: Optional[int] = ...
) -> tuple[Array, Array]: ...
def triu_indices_from(arr: ArrayLike, k: int = ...) -> tuple[Array, Array]: ...
def true_divide(x: ArrayLike, y: ArrayLike, /) -> Array: ...
def trunc(x: ArrayLike, /) -> Array: ...
uint: Any
uint16: Any
uint32: Any
uint4: Any
uint64: Any
uint8: Any
def union1d(
    ar1: ArrayLike,
    ar2: ArrayLike,
    *,
    size: Optional[int] = ...,
    fill_value: Optional[ArrayLike] = ...,
) -> Array: ...
class _UniqueAllResult(NamedTuple):
  values: Array
  indices: Array
  inverse_indices: Array
  counts: Array
class _UniqueCountsResult(NamedTuple):
    values: Array
    counts: Array
class _UniqueInverseResult(NamedTuple):
    values: Array
    inverse_indices: Array
def unique(ar: ArrayLike, return_index: builtins.bool = ..., return_inverse: builtins.bool = ...,
           return_counts: builtins.bool = ..., axis: Optional[int] = ...,
           *, equal_nan: builtins.bool = ..., size: Optional[int] = ...,
           fill_value: Optional[ArrayLike] = ...
): ...
def unique_all(x: ArrayLike, /) -> _UniqueAllResult: ...
def unique_counts(x: ArrayLike, /) -> _UniqueCountsResult: ...
def unique_inverse(x: ArrayLike, /) -> _UniqueInverseResult: ...
def unique_values(x: ArrayLike, /) -> Array: ...
def unpackbits(
    a: ArrayLike,
    axis: Optional[int] = ...,
    count: Optional[ArrayLike] = ...,
    bitorder: str = ...,
) -> Array: ...
def unravel_index(indices: ArrayLike, shape: Shape) -> tuple[Array, ...]: ...
unsignedinteger = _np.unsignedinteger
def unwrap(p: ArrayLike, discont: Optional[ArrayLike] = ...,
           axis: int = ..., period: ArrayLike = ...) -> Array: ...
def vander(
    x: ArrayLike, N: Optional[int] = ..., increasing: builtins.bool = ...
) -> Array: ...
def var(a: ArrayLike, axis: _Axis = ..., dtype: DTypeLike = ...,
        out: None = ..., ddof: int = ..., keepdims: builtins.bool = ..., *,
        where: Optional[ArrayLike] = ...) -> Array: ...
def vdot(
    a: ArrayLike, b: ArrayLike, *, precision: PrecisionLike = ...,
    preferred_element_type: Optional[DTypeLike] = ...) -> Array: ...
def vecdot(x1: ArrayLike, x2: ArrayLike, /, *, axis: int = ...,
           precision: PrecisionLike = ...,
           preferred_element_type: Optional[DTypeLike] = ...) -> Array: ...
def vsplit(
    ary: ArrayLike, indices_or_sections: Union[int, ArrayLike]
) -> list[Array]: ...

def vstack(tup: Union[_np.ndarray, Array, Sequence[ArrayLike]],
           dtype: Optional[DTypeLike] = ...) -> Array: ...

@overload
def where(condition: ArrayLike, x: Literal[None] = ..., y: Literal[None] = ...,
          /, *, size: Optional[int] = ...,
          fill_value: Union[None, ArrayLike, tuple[ArrayLike, ...]] = ...
          ) -> tuple[Array, ...]: ...

@overload
def where(condition: ArrayLike, x: ArrayLike, y: ArrayLike, /, *,
          size: Optional[int] = ...,
          fill_value: Union[None, ArrayLike, tuple[ArrayLike, ...]] = ...
          ) -> Array: ...

@overload
def where(condition: ArrayLike, x: Optional[ArrayLike] = ...,
          y: Optional[ArrayLike] = ..., /, *, size: Optional[int] = ...,
          fill_value: Union[None, ArrayLike, tuple[ArrayLike, ...]] = ...
          ) -> Union[Array, tuple[Array, ...]]: ...

def zeros(shape: Any, dtype: Optional[DTypeLike] = ...,
          device: Optional[Union[_Device, _Sharding]] = ...) -> Array: ...
def zeros_like(a: Union[ArrayLike, DuckTypedArray],
               dtype: Optional[DTypeLike] = ...,
               shape: Any = ..., *,
               device: Optional[Union[_Device, _Sharding]] = ...) -> Array: ...

def vectorize(pyfunc, *, excluded = ..., signature = ...) -> Callable: ...