from typing import Self, Sequence

from lucid.types import _Scalar, _ArrayOrScalar, _ShapeLike, _ArrayLikeInt


class _TensorOps:
    def __add__(self, other: Self | _ArrayOrScalar) -> Self: ...

    def __radd__(self, other: Self | _ArrayOrScalar) -> Self: ...

    def __sub__(self, other: Self | _ArrayOrScalar) -> Self: ...

    def __rsub__(self, other: Self | _ArrayOrScalar) -> Self: ...

    def __mul__(self, other: Self | _ArrayOrScalar) -> Self: ...

    def __rmul__(self, other: Self | _ArrayOrScalar) -> Self: ...

    def __truediv__(self, other: Self | _ArrayOrScalar) -> Self: ...

    def __rtruediv__(self, other: Self | _ArrayOrScalar) -> Self: ...

    def __matmul__(self, other: Self | _ArrayOrScalar) -> Self: ...

    def __eq__(self, other: Self | _ArrayOrScalar) -> Self: ...

    def __ne__(self, other: Self | _ArrayOrScalar) -> Self: ...

    def __gt__(self, other: Self | _ArrayOrScalar) -> Self: ...

    def __ge__(self, other: Self | _ArrayOrScalar) -> Self: ...

    def __lt__(self, other: Self | _ArrayOrScalar) -> Self: ...

    def __le__(self, other: Self | _ArrayOrScalar) -> Self: ...

    def __pow__(self, _: _Scalar) -> Self: ...

    def __neg__(self) -> Self: ...

    def __and__(self, other: Self | _ArrayOrScalar) -> Self: ...

    def __rand__(self, other: Self | _ArrayOrScalar) -> Self: ...

    def __or__(self, other: Self | _ArrayOrScalar) -> Self: ...

    def __ror__(self, other: Self | _ArrayOrScalar) -> Self: ...

    @property
    def T(self) -> Self: ...

    @property
    def mT(self) -> Self: ...

    def dot(self, other: Self | _ArrayOrScalar) -> Self: ...

    def matmul(self, other: Self) -> Self: ...

    def sum(
        self, axis: int | tuple[int] | None = None, keepdims: bool = False
    ) -> Self: ...

    def mean(
        self, axis: int | tuple[int] | None = None, keepdims: bool = False
    ) -> Self: ...

    def var(
        self, axis: int | tuple[int] | None = None, keepdims: bool = False
    ) -> Self: ...

    def clip(
        self, min_value: _Scalar | None = None, max_value: _Scalar | None = None
    ) -> Self: ...

    def reshape(self, *shape: int) -> Self: ...

    def transpose(self, axes: Sequence[int]) -> Self: ...

    def squeeze(self, axis: _ShapeLike | None = None) -> Self: ...

    def unsqueeze(self, axis: _ShapeLike) -> Self: ...

    def ravel(self) -> Self: ...

    def pad(self, pad_width: _ArrayLikeInt) -> Self: ...

    def repeat(self, repeats: int | Sequence[int], axis: int | None = None) -> Self: ...

    def tile(self, reps: int | Sequence[int]) -> Self: ...

    def flatten(self, start_axis: int = 0, end_axis: int = -1) -> Self: ...

    def split(
        self, size_or_sections: int | Sequence[int], axis: int = 0
    ) -> tuple[Self, ...]: ...

    def tril(self, diagonal: int = 0) -> Self: ...

    def triu(self, diagonal: int = 0) -> Self: ...

    def broadcast_to(self, shape: _ShapeLike) -> Self: ...

    def chunk(self, chunks: int, axis: int = 0) -> tuple[Self, ...]: ...

    def swapaxes(self, axis1: int, axis2: int) -> Self: ...

    def masked_fill(self, mask: Self, value: _Scalar) -> Self: ...

    def roll(
        self, shifts: int | tuple[int, ...], axis: int | tuple[int, ...] | None = None
    ) -> Self: ...

    def round(self, decimals: int = 0) -> Self: ...

    def floor(self) -> Self: ...

    def ceil(self) -> Self: ...

    def unbind(self, axis: int = 0) -> tuple[Self, ...]: ...

    def sort(self, axis: int = -1, descending: bool = False) -> tuple[Self, Self]: ...

    def unique(self, sorted: bool = True, axis: int | None = None) -> Self: ...

    def nonzero(self) -> Self: ...

    def diagonal(self, offset: int = 0, axis1: int = 0, axis2: int = 1) -> Self: ...
