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