tensor.py

import numpy as np
from numpy.typing import ArrayLike
from typing import Any, Callable
from abc import ABC, abstractmethod
from functools import partial
from .backend import BackendInterface, NumpyBackend, CupyBackend, NumbaBackend


class Tensor(object):
    __slots__ = ['_backend', 'data', 'gradient', 'requireGradient', 'gradientFunc', 'batched']

    __backend__ = NumpyBackend()

    def __init__(self, data: Any,
                 gradient: Any = None,
                 gradientFunc: Callable = None,
                 requireGradient: bool = False,
                 batched: bool = True) -> None:

        self._backend = Tensor.__backend__

        #if isinstance(data, (list | np.ndarray)):
        #    data = self._backend.array(data)
        #elif isinstance(data, (int, float)):
        #    data = self._backend.array([data])
        #elif isinstance(data, self.__class__):
        #    gradient = data.gradient if gradient is None else gradient
        #    gradientFunc = data.gradientFunc if gradientFunc is None else gradientFunc
        #    requireGradient = data.requireGradient if requireGradient is False else requireGradient
        #    data = data.data

        #if len(data.shape) == 1:
        #    data = self._backend.reshape(data, (1, *data.shape))

        #if gradient is None and requireGradient:
        #    # If gradient is not provided and it's required, initialize it as None
        #    gradient = self._backend.zeros_like(data)
        #elif isinstance(gradient, (list, int, float)):
        #    gradient = self._backend.array(gradient)

        # Checking if the shapes are the same
        #if gradient is not None:
        #    assert data.shape == gradient.shape, "value and gradient must have the same shape"

        self.data = data
        self.gradient = gradient
        self.requireGradient = requireGradient
        self.gradientFunc = gradientFunc
        self.batched = batched

    def zeroGradient(self) -> None:
        """In-place operation for nulling the gradient"""
        if self.requireGradient:
            self.gradient = self._backend.zeros_like(self.data)
        else:
            raise AttributeError("this tensor is not differentiable")

    def backward(self, gradient=None):
        """
        Compute the gradients recursively by applying the chain rule.
        """
        if gradient is None:
            gradient = self._backend.ones_like(self.data)

        if not self.requireGradient:
            return

        # If grad_fn is not set, this is probably the starting point for backpropagation,
        # so we don't need to compute further backward.
        if self.gradientFunc is None:
            return

        # Accumulate gradients instead of overwriting.
        self.gradient += gradient
        # Compute the local gradients using grad_fn
        self.gradientFunc.backward(self.gradient)

    def __repr__(self) -> str:
        """String representation."""
        dataTitle = 'data:\n'
        gradientTitle = 'gradient:\n'
        dataStr = str(self.data)
        gradientStr = str(self.gradient)
        if self.requireGradient is True:
            return dataTitle + dataStr + '\n' + gradientTitle + gradientStr
        else:
            return dataTitle + dataStr

    def copy(self) -> 'Tensor':
        data = self._backend.copy(self.data)
        gradient = self._backend.copy(self.gradient)
        return self.__class__(data, gradient, gradientFunc=self.gradientFunc, requireGradient=self.requireGradient)

    @property
    def strides(self) -> tuple:
        return self.data.strides

    def __len__(self) -> int:
        """Return the length of the value."""
        return len(self.data)

    @property
    def shape(self) -> tuple:
        """Return the shape of the value."""
        return self.data.shape

    @property
    def ndim(self) -> tuple:
        """Return the ndim of the value."""
        return self.data.ndim

    def reshape(self, newShape) -> 'Tensor':
        return Reshape()(self, newShape)

    def transpose(self) -> 'Tensor':
        return Transpose()(self)

    def T(self) -> 'Tensor':
        return Transpose()(self)

    def tolist(self) -> tuple[list, list] | list:
        if self.requireGradient is True:
            return self.data.tolist(), self.gradient.tolist()
        else:
            return self.data.tolist()

    @classmethod
    def setBackend(cls, backend: BackendInterface) -> None:
        cls.__backend__ = backend

    def __getitem__(self, index):
        """Get an item by index."""
        if self.requireGradient is True and self.gradient:
            return self.__class__(data=self.data[index], gradient=self.gradient[index], requireGradient=True, gradientFunc=self.gradientFunc)
        elif self.requireGradient is True:
            return self.__class__(data=self.data[index], requireGradient=True, gradientFunc=self.gradientFunc)
        else:
            return self.__class__(data=self.data[index], requireGradient=False)

    def __setitem__(self, index, value) -> None:
        """Set the value of an item by index."""
        if isinstance(value, self.__class__):
            self.data[index] = value.data
            if self.requireGradient is True and self.gradient:
                self.gradient[index] = value.gradient
                self.requireGradient = True
        else:
            self.data[index] = value
            self.gradient[index] = 0

    def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
        if method == '__call__':
            operation = ufuncMap.get(ufunc)
            if operation is not None:
                return operation()(*inputs, **kwargs)
        raise NotImplementedError(f'{ufunc} is not implemented yet')

    def __array_function__(self, func, types, args, kwargs):
        operation = funcMap.get(func)
        if operation is not None:
            return operation()(*args, **kwargs)
        raise NotImplementedError(f'{func} is not implemented yet')

    def __add__(self, other: ArrayLike) -> 'Tensor':
        return addForward(self, other)

    def __radd__(self, other: ArrayLike) -> 'Tensor':
        return addForward(other, self)

    def __iadd__(self, other: ArrayLike) -> 'Tensor':
        result = addForward(self, other)
        self.data = result.data
        self.gradient = result.gradient
        self.requireGradient = result.requireGradient
        return self

    def __sub__(self, other: ArrayLike) -> 'Tensor':
        return subtractForward(self, other)

    def __rsub__(self, other: ArrayLike) -> 'Tensor':
        return subtractForward(other, self)

    def __isub__(self, other: ArrayLike) -> 'Tensor':
        result = subtractForward(self, other)
        self.data = result.data
        self.gradient = result.gradient
        self.requireGradient = result.requireGradient
        return self

    def __mul__(self, other: ArrayLike) -> 'Tensor':
        return Multiply()(self, other)

    def __rmul__(self, other: ArrayLike) -> 'Tensor':
        return Multiply()(other, self)

    def __imul__(self, other: ArrayLike) -> 'Tensor':
        result = Multiply(self, other)
        self.data = result.data
        self.gradient = result.gradient
        self.requireGradient = result.requireGradient
        return self

    def __truediv__(self, other: ArrayLike) -> 'Tensor':
        return Divide()(self, other)

    def __rtruediv__(self, other: ArrayLike) -> 'Tensor':
        return Divide()(other, self)

    def __itruediv__(self, other: ArrayLike) -> 'Tensor':
        result = Divide(self, other)
        self.data = result.data
        self.gradient = result.gradient
        self.requireGradient = result.requireGradient
        return self

    def __matmul__(self, other: ArrayLike) -> 'Tensor':
        return Matmul()(self, other)

    def __rmatmul__(self, other: ArrayLike) -> 'Tensor':
        return Matmul()(other, self)

    def __imatmul__(self, other: ArrayLike) -> 'Tensor':
        result = Matmul(self, other)
        self.data = result.data
        self.gradient = result.gradient
        self.requireGradient = result.requireGradient
        return self

    def __pow__(self, other: ArrayLike) -> 'Tensor':
        return Power()(self, other)

    def __rpow__(self, other: ArrayLike) -> 'Tensor':
        return Power()(other, self)

    def __ipow__(self, other: ArrayLike) -> 'Tensor':
        result = Power(self, other)
        self.data = result.data
        self.gradient = result.gradient
        self.requireGradient = result.requireGradient
        return self

    def __abs__(self) -> 'Tensor':
        return Abs()(self)

    def __pos__(self) -> 'Tensor':
        return Positive()(self)

    def __neg__(self) -> 'Tensor':
        return Negative()(self)

    def __eq__(self, other) -> bool:
        """Equality comparison."""
        return Equal()(self, other)

    def __gt__(self, other) -> bool:
        """Greater than comparison."""
        return Greater()(self, other)

    def __ge__(self, other) -> bool:
        """Greater than or equal to comparison."""
        return GreaterEqual()(self, other)

    def __lt__(self, other) -> bool:
        """Less than comparison."""
        return Less()(self, other)

    def __le__(self, other) -> bool:
        """Less than or equal to comparison."""
        return LessEqual()(self, other)


def checkTensor(tensor: Tensor) -> Tensor:
    if isinstance(tensor, Tensor):
        return tensor
    return Tensor(tensor)


#
# Operations
#


class Operation(ABC):
    __slots__ = ['name', 'operationID', 'backend']
    id = 0
    __backend__ = Tensor.__backend__

    def __init__(self) -> None:
        self.name = self.__class__.__name__
        self.operationID = Operation.id
        self.backend = Operation.__backend__
        Operation.id += 1

    @abstractmethod
    def forward(self, *args, **kwargs) -> Tensor:
        raise NotImplementedError

    @abstractmethod
    def backward(self, gradient: np.ndarray) -> None:
        raise NotImplementedError

    def __call__(self, *args, **kwargs):
        return self.forward(*args, **kwargs)

    def __repr__(self) -> str:
        return f'{self.name}, {self.operationID}'

    def at(self, tensor: Tensor, indices, value) -> None:
         # not ready for use yet
         tensor = self.forward(indices, value)


class TwoTensors(Operation):
    __slots__ = ['tensor1', 'tensor2']

    def __init__(self) -> None:
        super().__init__()
        self.tensor1 = None
        self.tensor2 = None
        self.tensor1BroadcastAxis = None
        self.tensor2BroadcastAxis = None

    def getbroadcastAxid(self, data, gradient) -> None:
        # Store old shapes
        tensorShape = np.array(data.shape)

        # Get new shape
        gradientShape = np.array(gradient.shape)

        # Prepend ones to the shape of the smaller array
        if len(tensorShape) < len(gradientShape):
            tensorShape = np.pad(tensorShape, (len(gradientShape) - len(tensorShape), 0), mode='constant', constant_values=1)
        elif len(tensorShape) > len(gradientShape):
            gradientShape = np.pad(gradientShape, (len(tensorShape) - len(gradientShape), 0), mode='constant', constant_values=1)

        # Find broadcasted axes
        tensorBroadcastAxis = np.where(tensorShape != gradientShape)[0]

        # Change broadcastAxis variables to None if they're empty
        if tensorBroadcastAxis.size == 0:
            tensorBroadcastAxis = None

        return tensorBroadcastAxis

    def forward(self, tensor1: Tensor, tensor2: Tensor, *args, **kwargs) -> Tensor:
        if not isinstance(tensor1, Tensor):
            tensor1 = Tensor(tensor1)
        if not isinstance(tensor2, Tensor):
            tensor2 = Tensor(tensor2)

        requireGradient = tensor1.requireGradient or tensor2.requireGradient
        if requireGradient:
            self.tensor1 = tensor1
            self.tensor2 = tensor2

        data = self._operation(tensor1.data, tensor2.data, *args, **kwargs)

        return Tensor(data, requireGradient=requireGradient, gradientFunc=self)

    def backward(self, gradient: np.ndarray) -> None:
        if self.tensor1 and self.tensor1.requireGradient:
            gradientForTensor1 = self.backend.copy(gradient)

            tensorBroadcastAxis = self.getbroadcastAxid(self.tensor1, gradientForTensor1)
            if tensorBroadcastAxis is not None:
                gradientForTensor1 = np.sum(gradientForTensor1, axis=tuple(tensorBroadcastAxis), keepdims=True)

            self.tensor1.gradient = self._derivativeD1(gradientForTensor1)
            if self.tensor1.gradientFunc:
                self.tensor1.gradientFunc.backward(self.tensor1.gradient)

        if self.tensor2 and self.tensor2.requireGradient:
            gradientForTensor2 = self.backend.copy(gradient)

            tensorBroadcastAxis = self.getbroadcastAxid(self.tensor2, gradientForTensor2)
            if tensorBroadcastAxis is not None:
                gradientForTensor2 = np.sum(gradientForTensor2, axis=tuple(tensorBroadcastAxis), keepdims=True)

            self.tensor2.gradient = self._derivativeD2(gradientForTensor2)
            if self.tensor2.gradientFunc:
                self.tensor2.gradientFunc.backward(self.tensor2.gradient)

    @abstractmethod
    def _operation(self, data1: np.ndarray, data2: np.ndarray, *args, **kwargs) -> np.ndarray:
        raise NotImplementedError

    @abstractmethod
    def _derivativeD1(self, gradient: np.ndarray, *args, **kwargs) -> np.ndarray:
        raise NotImplementedError

    @abstractmethod
    def _derivativeD2(self, gradient: np.ndarray, *args, **kwargs) -> np.ndarray:
        raise NotImplementedError


class OneTensor(Operation):
    __slots__ = ['tensor']

    def __init__(self) -> None:
        super().__init__()
        self.tensor = None

    def forward(self, tensor: Tensor, *args, **kwargs) -> Tensor:
        if not isinstance(tensor, Tensor):
            tensor = Tensor(tensor)

        requireGradient = tensor.requireGradient
        if requireGradient:
            self.tensor = tensor

        data = self._operation(tensor.data, *args, **kwargs)

        return Tensor(data, requireGradient=requireGradient, gradientFunc=self)

    def backward(self, gradient: np.ndarray) -> None:
        if self.tensor and self.tensor.requireGradient:
            self.tensor.gradient = self._derivative(gradient)
            if self.tensor.gradientFunc:
                self.tensor.gradientFunc.backward(self.tensor.gradient)

    @abstractmethod
    def _operation(self, data: np.ndarray, *args, **kwargs) -> np.ndarray:
        raise NotImplementedError

    @abstractmethod
    def _derivative(self, gradient: np.ndarray, *args, **kwargs) -> np.ndarray:
        raise NotImplementedError


#
# Two Tensors
#


def getbroadcastAxid(data, gradient) -> None:
    # Store old shapes
    tensorShape = np.array(data.shape)

    # Get new shape
    gradientShape = np.array(gradient.shape)

    # Prepend ones to the shape of the smaller array
    if len(tensorShape) < len(gradientShape):
        tensorShape = np.pad(tensorShape, (len(gradientShape) - len(tensorShape), 0), mode='constant', constant_values=1)
    elif len(tensorShape) > len(gradientShape):
        gradientShape = np.pad(gradientShape, (len(tensorShape) - len(gradientShape), 0), mode='constant', constant_values=1)

    # Find broadcasted axes
    tensorBroadcastAxis = np.where(tensorShape != gradientShape)[0]

    # Change broadcastAxis variables to None if they're empty
    if tensorBroadcastAxis.size == 0:
        tensorBroadcastAxis = None

    return tensorBroadcastAxis


def addForward(tensor1: Tensor, tensor2: Tensor, *args, **kwargs) -> Tensor:
    if not isinstance(tensor1, Tensor):
        tensor1 = Tensor(tensor1)
    if not isinstance(tensor2, Tensor):
        tensor2 = Tensor(tensor2)

    data = np.add(tensor1.data, tensor2.data, *args, **kwargs)

    requireGradient = tensor1.requireGradient or tensor2.requireGradient
    if requireGradient:
        gradfunc = partial(addBackward, tensor1, tensor2)
        return Tensor(data, requireGradient=requireGradient, gradientFunc=gradfunc)

    return Tensor(data, requireGradient=requireGradient, gradientFunc=None)


def addBackward(tensor1: Tensor, tensor2: Tensor, gradient: np.ndarray, *args, **kwargs) -> np.ndarray:
    if tensor1 and tensor1.requireGradient:
        gradientForTensor1 = np.copy(gradient)

        tensorBroadcastAxis = getbroadcastAxid(tensor1, gradientForTensor1)
        if tensorBroadcastAxis is not None:
            gradientForTensor1 = np.sum(gradientForTensor1, axis=tuple(tensorBroadcastAxis), keepdims=True)

        tensor1.gradient = np.add(tensor1.gradient, gradient)
        if tensor1.gradientFunc:
            tensor1.gradientFunc.backward(tensor1.gradient)

    if tensor2 and tensor2.requireGradient:
        gradientForTensor2 = np.copy(gradient)

        tensorBroadcastAxis = getbroadcastAxid(tensor2, gradientForTensor2)
        if tensorBroadcastAxis is not None:
            gradientForTensor2 = np.sum(gradientForTensor2, axis=tuple(tensorBroadcastAxis), keepdims=True)

        tensor2.gradient = np.add(tensor2.gradient, gradient)
        if tensor2.gradientFunc:
            tensor2.gradientFunc.backward(tensor2.gradient)


class Add(TwoTensors):
    def __init__(self) -> None:
        super().__init__()

    def _operation(self, data1: np.ndarray, data2: np.ndarray, *args, **kwargs) -> np.ndarray:
        return self.backend.add(data1, data2, *args, **kwargs)

    def _derivativeD1(self, gradient: np.ndarray, *args, **kwargs) -> np.ndarray:
        return self.backend.add(self.tensor1.gradient, gradient)

    def _derivativeD2(self, gradient: np.ndarray, *args, **kwargs) -> np.ndarray:
        return self.backend.add(self.tensor2.gradient, gradient)


def subtractForward(tensor1: Tensor, tensor2: Tensor, *args, **kwargs) -> Tensor:
    if not isinstance(tensor1, Tensor):
        tensor1 = Tensor(tensor1)
    if not isinstance(tensor2, Tensor):
        tensor2 = Tensor(tensor2)

    data = np.subtract(tensor1.data, tensor2.data, *args, **kwargs)

    requireGradient = tensor1.requireGradient or tensor2.requireGradient
    if requireGradient:
        gradfunc = partial(addBackward, tensor1, tensor2)
        return Tensor(data, requireGradient=requireGradient, gradientFunc=gradfunc)

    return Tensor(data, requireGradient=requireGradient, gradientFunc=None)


def subtractBackward(tensor1: Tensor, tensor2: Tensor, gradient: np.ndarray, *args, **kwargs) -> np.ndarray:
    if tensor1 and tensor1.requireGradient:
        gradientForTensor1 = np.copy(gradient)

        tensorBroadcastAxis = getbroadcastAxid(tensor1, gradientForTensor1)
        if tensorBroadcastAxis is not None:
            gradientForTensor1 = np.sum(gradientForTensor1, axis=tuple(tensorBroadcastAxis), keepdims=True)

        tensor1.gradient = np.add(tensor1.gradient, gradient)
        if tensor1.gradientFunc:
            tensor1.gradientFunc.backward(tensor1.gradient)

    if tensor2 and tensor2.requireGradient:
        gradientForTensor2 = np.copy(gradient)

        tensorBroadcastAxis = getbroadcastAxid(tensor2, gradientForTensor2)
        if tensorBroadcastAxis is not None:
            gradientForTensor2 = np.sum(gradientForTensor2, axis=tuple(tensorBroadcastAxis), keepdims=True)

        tensor2.gradient = np.subtract(tensor2.gradient, gradient)
        if tensor2.gradientFunc:
            tensor2.gradientFunc.backward(tensor2.gradient)


class Subtract(TwoTensors):
    def __init__(self) -> None:
        super().__init__()

    def _operation(self, data1: np.ndarray, data2: np.ndarray, *args, **kwargs) -> np.ndarray:
        return self.backend.subtract(data1, data2, *args, **kwargs)

    def _derivativeD1(self, gradient: np.ndarray, *args, **kwargs) -> np.ndarray:
        return self.backend.add(self.tensor1.gradient, gradient)

    def _derivativeD2(self, gradient: np.ndarray, *args, **kwargs) -> np.ndarray:
        return self.backend.subtract(self.tensor2.gradient, gradient)


class Multiply(TwoTensors):
    def __init__(self) -> None:
        super().__init__()

    def _operation(self, data1: np.ndarray, data2: np.ndarray, *args, **kwargs) -> np.ndarray:
        return self.backend.multiply(data1, data2, *args, **kwargs)

    def _derivativeD1(self, gradient: np.ndarray, *args, **kwargs) -> np.ndarray:
        return self.backend.add(self.tensor1.gradient, self.backend.multiply(self.tensor2.data, gradient))

    def _derivativeD2(self, gradient: np.ndarray, *args, **kwargs) -> np.ndarray:
        return self.backend.add(self.tensor2.gradient, self.backend.multiply(self.tensor1.data, gradient))


class Divide(TwoTensors):
    def __init__(self) -> None:
        super().__init__()

    def _operation(self, data1: np.ndarray, data2: np.ndarray, *args, **kwargs) -> np.ndarray:
        return self.backend.divide(data1, data2, *args, **kwargs)

    def _derivativeD1(self, gradient: np.ndarray, *args, **kwargs) -> np.ndarray:
        return self.backend.add(self.tensor1.gradient, self.backend.divide(gradient, self.tensor2.data))

    def _derivativeD2(self, gradient: np.ndarray, *args, **kwargs) -> np.ndarray:
        return self.backend.subtract(self.tensor2.gradient, self.backend.divide(self.backend.multiply(self.tensor1.data, gradient), self.backend.power(self.tensor2.data, 2)))


class Matmul(TwoTensors):
    def __init__(self) -> None:
        super().__init__()

    def _operation(self, data1: np.ndarray, data2: np.ndarray, *args, **kwargs) -> np.ndarray:
        return self.backend.matmul(data1, data2, *args, **kwargs)

    # Update the backward pass to handle batch dimension
    def _derivativeD1(self, gradient: np.ndarray, *args, **kwargs) -> np.ndarray:
        if len(self.tensor1.data.shape) > 2 or len(self.tensor2.data.shape) > 2:
            return self.backend.add(self.tensor1.gradient, self.backend.matmul(gradient, self.backend.transpose(self.tensor2.data, axes=(0, 2, 1))))
        else:
            return self.backend.add(self.tensor1.gradient, self.backend.matmul(gradient, self.backend.transpose(self.tensor2.data)))

    def _derivativeD2(self, gradient: np.ndarray, *args, **kwargs) -> np.ndarray:
        if len(self.tensor1.data.shape) > 2 or len(self.tensor2.data.shape) > 2:
            return self.backend.add(self.tensor2.gradient, self.backend.matmul(self.backend.transpose(self.tensor1.data, axes=(0, 2, 1)), gradient))
        else:
            return self.backend.add(self.tensor2.gradient, self.backend.matmul(self.backend.transpose(self.tensor1.data), gradient))

    def backward(self, gradient: np.ndarray) -> None:
        if self.tensor1 and self.tensor1.requireGradient:
            gradientForTensor1 = self.backend.copy(gradient)

            self.tensor1.gradient = self._derivativeD1(gradientForTensor1)
            if self.tensor1.gradientFunc:
                self.tensor1.gradientFunc.backward(self.tensor1.gradient)

        if self.tensor2 and self.tensor2.requireGradient:
            gradientForTensor2 = self.backend.copy(gradient)

            self.tensor2.gradient = self._derivativeD2(gradientForTensor2)
            if self.tensor2.gradientFunc:
                self.tensor2.gradientFunc.backward(self.tensor2.gradient)


class Dot(TwoTensors):
    def __init__(self) -> None:
        super().__init__()

    def _operation(self, data1: np.ndarray, data2: np.ndarray, *args, **kwargs) -> np.ndarray:
        return self.backend.dot(data1, data2)

    def _derivativeD1(self, gradient: np.ndarray, *args, **kwargs) -> np.ndarray:
        return self.backend.multiply(self.tensor2.data, gradient)

    def _derivativeD2(self, gradient: np.ndarray, *args, **kwargs) -> np.ndarray:
        return self.backend.multiply(self.tensor1.data, gradient)


class Power(TwoTensors):
    def __init__(self) -> None:
        super().__init__()

    def _operation(self, data1: np.ndarray, data2: np.ndarray, *args, **kwargs) -> np.ndarray:
        return self.backend.power(data1, data2)

    def _derivativeD1(self, gradient: np.ndarray, *args, **kwargs) -> np.ndarray:
        return self.backend.add(self.tensor1.gradient, self.backend.multiply(self.backend.multiply(self.tensor2.data, self.backend.power(self.tensor1.data, (self.backend.subtract(self.tensor2.data, 1)))), gradient))

    def _derivativeD2(self, gradient: np.ndarray, *args, **kwargs) -> np.ndarray:
        return self.backend.add(self.tensor2.gradient, self.backend.multiply(self.backend.multiply(self.backend.log(self.tensor1.data), self.backend.power(self.tensor1.data, self.tensor2.data)), gradient))


#
# Single Tensor
#


class Square(OneTensor):
    def __init__(self) -> None:
        super().__init__()

    def _operation(self, data: np.ndarray, *args, **kwargs) -> np.ndarray:
        return self.backend.square(data, *args, **kwargs)

    def _derivative(self, gradient: np.ndarray, *args, **kwargs) -> np.ndarray:
        return self.backend.multiply(self.backend.multiply(self.tensor.data, 2.0), gradient)


class Sqrt(OneTensor):
    def __init__(self) -> None:
        super().__init__()

    def _operation(self, data: np.ndarray, *args, **kwargs) -> np.ndarray:
        return self.backend.sqrt(data, *args, **kwargs)

    def _derivative(self, gradient: np.ndarray, *args, **kwargs) -> np.ndarray:
        return self.backend.multiply(self.backend.divide(0.5, self.backend.sqrt(self.tensor.data)), gradient)


class Log(OneTensor):
    def __init__(self) -> None:
        super().__init__()

    def _operation(self, data: np.ndarray, *args, **kwargs) -> np.ndarray:
        return self.backend.log(data, *args, **kwargs)

    def _derivative(self, gradient: np.ndarray, *args, **kwargs) -> np.ndarray:
        return self.backend.add(self.tensor.gradient, self.backend.multiply((self.backend.divide(1, self.tensor.data)), gradient))


class Exp(OneTensor):
    __slots__ = ['data']

    def __init__(self) -> None:
        super().__init__()
        self.data = None

    def _operation(self, data: np.ndarray, *args, **kwargs) -> np.ndarray:
        self.data = self.backend.exp(data, *args, **kwargs)
        return self.data

    def _derivative(self, gradient: np.ndarray, *args, **kwargs) -> np.ndarray:
        return self.data * gradient


class Sin(OneTensor):
    def __init__(self) -> None:
        super().__init__()

    def _operation(self, data: np.ndarray, *args, **kwargs) -> np.ndarray:
        return self.backend.sin(data, *args, **kwargs)

    def _derivative(self, gradient: np.ndarray, *args, **kwargs) -> np.ndarray:
        return self.backend.multiply(self.backend.cos(self.tensor.data), gradient)


class Cos(OneTensor):
    def __init__(self) -> None:
        super().__init__()

    def _operation(self, data: np.ndarray, *args, **kwargs) -> np.ndarray:
        return self.backend.cos(data, *args, **kwargs)

    def _derivative(self, gradient: np.ndarray, *args, **kwargs) -> np.ndarray:
        return self.backend.negative(self.backend.multiply(self.backend.sin(self.tensor.data), gradient))


class Tan(OneTensor):
    def __init__(self) -> None:
        super().__init__()

    def _operation(self, data: np.ndarray, *args, **kwargs) -> np.ndarray:
        return self.backend.tan(data, *args, **kwargs)

    def _derivative(self, gradient: np.ndarray, *args, **kwargs) -> np.ndarray:
        return self.backend.multiply((self.backend.divide(1, self.backend.power(np.cos(self.tensor.data), 2))), gradient)


class Sinh(OneTensor):
    def __init__(self) -> None:
        super().__init__()

    def _operation(self, data: np.ndarray, *args, **kwargs) -> np.ndarray:
        return self.backend.sinh(data, *args, **kwargs)

    def _derivative(self, gradient: np.ndarray, *args, **kwargs) -> np.ndarray:
        return self.backend.multiply(self.backend.cosh(self.tensor.data), gradient)


class Cosh(OneTensor):
    def __init__(self) -> None:
        super().__init__()

    def _operation(self, data: np.ndarray, *args, **kwargs) -> np.ndarray:
        return self.backend.cosh(data, *args, **kwargs)

    def _derivative(self, gradient: np.ndarray, *args, **kwargs) -> np.ndarray:
        return self.backend.multiply(self.backend.sinh(self.tensor.data), gradient)


class Tanh(OneTensor):
    def __init__(self) -> None:
        super().__init__()

    def _operation(self, data: np.ndarray, *args, **kwargs) -> np.ndarray:
        return self.backend.tanh(data, *args, **kwargs)

    def _derivative(self, gradient: np.ndarray, *args, **kwargs) -> np.ndarray:
        return self.backend.multiply((self.backend.divide(1, self.backend.power(np.cosh(self.tensor.data), 2))), gradient)


class Abs(OneTensor):
    def __init__(self) -> None:
        super().__init__()

    def _operation(self, data: np.ndarray, *args, **kwargs) -> np.ndarray:
        return self.backend.abs(data, *args, **kwargs)

    def _derivative(self, gradient: np.ndarray, *args, **kwargs) -> np.ndarray:
        return self.backend.multiply(self.backend.sign(self.tensor.data), gradient)


#
# Signs
#


class Sign(OneTensor):
    def __init__(self) -> None:
        super().__init__()

    def _operation(self, data: np.ndarray, *args, **kwargs) -> np.ndarray:
        return self.backend.sign(data, *args, **kwargs)

    def _derivative(self, gradient: np.ndarray, *args, **kwargs) -> np.ndarray:
        return self.backend.multiply(self.backend.sign(self.tensor.data), gradient)


class Positive(OneTensor):
    def __init__(self) -> None:
        super().__init__()

    def _operation(self, data: np.ndarray, *args, **kwargs) -> np.ndarray:
        return self.backend.positive(data, *args, **kwargs)

    def _derivative(self, gradient: np.ndarray, *args, **kwargs) -> np.ndarray:
        return self.backend.multiply(self.backend.positive(self.tensor.data), gradient)


class Negative(OneTensor):
    def __init__(self) -> None:
        super().__init__()

    def _operation(self, data: np.ndarray, *args, **kwargs) -> np.ndarray:
        return self.backend.negative(data, *args, **kwargs)

    def _derivative(self, gradient: np.ndarray, *args, **kwargs) -> np.ndarray:
        return self.backend.multiply(self.backend.negative(self.tensor.data), gradient)


#
# Compare
#


class Equal(TwoTensors):
    __slots__ = ['bools']

    def __init__(self) -> None:
        super().__init__()
        self.bools = None

    def _operation(self, data1: np.ndarray, data2: np.ndarray, *args, **kwargs) -> np.ndarray:
        return self.backend.equal(data1, data2)

    def _derivative(self, gradient: np.ndarray, *args, **kwargs) -> np.ndarray:
        return self.backend.multiply(self.bools, gradient)


class NotEqual(TwoTensors):
    __slots__ = ['bools']

    def __init__(self) -> None:
        super().__init__()

    def _operation(self, data1: np.ndarray, data2: np.ndarray, *args, **kwargs) -> np.ndarray:
        self.bools = self.backend.not_equal(data1, data2)
        return self.bools

    def _derivative(self, gradient: np.ndarray, *args, **kwargs) -> np.ndarray:
        return self.backend.multiply(self.bools, gradient)


class Less(TwoTensors):
    __slots__ = ['bools']

    def __init__(self) -> None:
        super().__init__()

    def _operation(self, data1: np.ndarray, data2: np.ndarray, *args, **kwargs) -> np.ndarray:
        self.bools = self.backend.less(data1, data2)
        return self.bools

    def _derivative(self, gradient: np.ndarray, *args, **kwargs) -> np.ndarray:
        return self.backend.multiply(self.bools, gradient)


class LessEqual(TwoTensors):
    __slots__ = ['bools']

    def __init__(self) -> None:
        super().__init__()

    def _operation(self, data1: np.ndarray, data2: np.ndarray, *args, **kwargs) -> np.ndarray:
        self.bools = self.backend.less_equal(data1, data2)
        return self.bools

    def _derivative(self, gradient: np.ndarray, *args, **kwargs) -> np.ndarray:
        return self.backend.multiply(self.bools, gradient)


class Greater(TwoTensors):
    __slots__ = ['bools']

    def __init__(self) -> None:
        super().__init__()

    def _operation(self, data1: np.ndarray, data2: np.ndarray, *args, **kwargs) -> np.ndarray:
        self.bools = self.backend.greater(data1, data2)
        return self.bools

    def _derivative(self, gradient: np.ndarray, *args, **kwargs) -> np.ndarray:
        return self.backend.multiply(self.bools, gradient)


class GreaterEqual(TwoTensors):
    __slots__ = ['bools']

    def __init__(self) -> None:
        super().__init__()

    def _operation(self, data1: np.ndarray, data2: np.ndarray, *args, **kwargs) -> np.ndarray:
        self.bools = self.backend.greater_equal(data1, data2)
        return self.bools

    def _derivative(self, gradient: np.ndarray, *args, **kwargs) -> np.ndarray:
        return self.backend.multiply(self.bools, gradient)


#
# Shaping
#


class Flatten(OneTensor):
    def __init__(self) -> None:
        super().__init__()

    def _operation(self, data: np.ndarray, *args, **kwargs) -> np.ndarray:
        return self.backend.reshape(data, newshape=(-1))

    def _derivative(self, gradient: np.ndarray, *args, **kwargs) -> np.ndarray:
        return self.backend.reshape(gradient, newshape=self.tensor.shape)


class Reshape(OneTensor):
    def __init__(self) -> None:
        super().__init__()

    def _operation(self, data: np.ndarray, *args, **kwargs) -> np.ndarray:
        return self.backend.reshape(data, *args, **kwargs)

    def _derivative(self, gradient: np.ndarray, *args, **kwargs) -> np.ndarray:
        return self.backend.reshape(gradient, newshape=self.tensor.shape)


#
# Broadcasting
#


class Repeat(Operation):
    __slots__ = ['repeats', 'axis', 'tensor']

    def __init__(self) -> None:
        super().__init__()
        self.tensor = None

    def forward(self, tensor: Tensor, repeats: ArrayLike, axis: int = None) -> Tensor:
        tensor = checkTensor(tensor)
        self.repeats = repeats
        self.axis = axis

        requireGradient = tensor.requireGradient
        if requireGradient:
            self.tensor = tensor

        data = self.backend.repeat(tensor.data, repeats=self.repeats, axis=self.axis)

        return Tensor(data, requireGradient=requireGradient, gradientFunc=self)

    def backward(self, gradient) -> None:
        if self.tensor and self.tensor.requireGradient:
            if self.axis is None:
                sum_axis = tuple(range(gradient.ndim)[::-self.repeats])
                counts = np.prod(self.repeats)
            else:
                sum_axis = self.axis
                counts = self.repeats

            grad = self.backend.sum(gradient, axis=sum_axis, keepdims=True)
            grad = self.backend.divide(grad, counts)
            self.tensor.gradient = self.backend.broadcast_to(grad, self.tensor.shape)

            if self.tensor.gradientFunc:
                self.tensor.gradientFunc.backward(self.tensor.gradient)


class Tile(Operation):
    __slots__ = ['tensor', 'reps']

    def __init__(self) -> None:
        super().__init__()
        self.tensor = None

    def forward(self, tensor: Tensor, reps: ArrayLike) -> Tensor:
        tensor = checkTensor(tensor)
        self.reps = reps

        requireGradient = tensor.requireGradient
        if requireGradient:
            self.tensor = tensor

        data = self.backend.tile(tensor.data, reps=self.reps)