from __future__ import annotations

from collections.abc import Callable, Iterator, Sequence
from types import EllipsisType
from typing import ParamSpec, TypeVar

P = ParamSpec("P")
R = TypeVar("R")

def jit(
    func: Callable[P, R], *, static_argnames: Sequence[str] | None = None
) -> Callable[P, R]: ...
def vmap(
    func: Callable[P, R], *, in_axes: None | int | Sequence[int | None] = None
) -> Callable[P, R]: ...

class Array:
    def __sub__(self, other: Array | float) -> Array: ...
    def __rsub__(self, other: Array | float) -> Array: ...
    def __add__(self, other: Array | float) -> Array: ...
    def __radd__(self, other: Array | float) -> Array: ...
    def __getitem__(
        self,
        ind: int
        | None
        | slice
        | EllipsisType
        | Array
        | tuple[int | None | slice | EllipsisType | Array, ...],
    ) -> Array: ...
    def __pow__(self, exp: Array | float) -> Array: ...
    def __div__(self, other: Array) -> Array: ...
    def __truediv__(self, other: Array | float) -> Array: ...
    def __rtruediv__(self, other: Array | float) -> Array: ...
    def __neg__(self) -> Array: ...
    def __mul__(self, other: Array | float) -> Array: ...
    def __rmul__(self, other: Array | float) -> Array: ...
    def __lt__(self, other: Array | float) -> Array: ...
    def __gt__(self, other: Array | float) -> Array: ...
    def __matmul__(self, other: Array) -> Array: ...
    def __iter__(self) -> Iterator[Array]: ...
    @property
    def shape(self) -> tuple[int, ...]: ...
    @property
    def T(self) -> Array: ...
    @property
    def real(self) -> Array: ...
    def sum(self, axis: int | None = None) -> Array: ...
    def std(self, axis: int | None = None) -> Array: ...
    def reshape(self, *shape: int) -> Array: ...
    def flatten(self) -> Array: ...
