tensor.cpp

#include <iostream>
#include <vector>
#include <string>
#include <random>
#include "tensor.h"


namespace tensor {

    //static methods

    Tensor* Tensor::fromVector(const std::vector<int>& dims, bool isGrad) {
        Tensor* result = new Tensor(dims);
        result->isGrad = isGrad;
        return result;
    }

    Tensor* Tensor::fromVector(const std::vector<float>& data, const std::vector<int>& dims, bool isGrad) {
        Tensor* result= new Tensor(data, dims);
        result->isGrad = isGrad;
        if(isGrad)
        {
            result->grad.resize(data.size(), 0.0f);
        }
        return result;
    }

    Tensor* Tensor::fromVector(const Tensor& t, const std::vector<int>& dims, bool isGrad) {
        Tensor* result= new Tensor(dims);
        result->data = t.data;
        result->isGrad = isGrad;
        return result;
    }

    //Constructor and Destructor

    Tensor::Tensor() {
        isGrad = false;
    }

    Tensor::Tensor(const std::vector<int>& dims, bool isGrad) {
        shape = dims;
        strides.resize(shape.size());
        strides[shape.size() - 1] = 1;
        for (int i = shape.size() - 2; i >= 0; i--) {
            strides[i] = strides[i + 1] * shape[i + 1];
        }
        data.resize(strides[0] * shape[0]);
        this->isGrad = isGrad;
        if(isGrad)
        {
            grad.resize(data.size(), 0.0f);
        }
    }

    Tensor::Tensor(Tensor& other) {
        shape = other.shape;
        strides = other.strides;
        data = other.data;
        isGrad = other.isGrad;
        grad = other.grad;
        previous = other.previous;
        gradFn = other.gradFn; 
        isLeafNode(other.isLeaf);
        name = other.name;
    }

    Tensor::Tensor(const std::vector<float>& data, const std::vector<int>& dims, const std::vector<int>& strides, bool isGrad) {
        this->data = data;
        this->shape = dims;
        this->strides = strides;
        this->isGrad = isGrad;
        if(isGrad)
        {
            grad.resize(data.size(), 0.0f);
        }
    }

    Tensor::Tensor(const std::vector<float>& data, const std::vector<int>& dims,bool isGrad) {
        shape = dims;
        strides.resize(shape.size());
        strides[shape.size() - 1] = 1;
        for (int i = shape.size() - 2; i >= 0; i--) {
            strides[i] = strides[i + 1] * shape[i + 1];
        }
        this->data = data;
        this->isGrad = isGrad;
        if(isGrad)
        {
            grad.resize(data.size(), 0.0f);
        }
    }

    Tensor::Tensor(const std::vector<int>& dims, Tensor* other) : Tensor(dims)
    {
        addToPrevious(other);
    }

    Tensor::~Tensor() {
    }
    
    //Normal Distribution
    Tensor* Tensor::normal(float mean, float std, const std::vector<int>& dims, bool isGrad) {
        static std::random_device rd;
        static std::mt19937 gen(rd());
        std::normal_distribution<float> distribution(mean, std);
        
        Tensor* result = new Tensor(dims, isGrad);
        for (int i = 0; i < result->data.size(); i++) {
            result->data[i] = distribution(gen);
        }
        return result;
    }

    //MatMul

    Tensor* Tensor::matmul2D_(Tensor* other) 
        {
            std::vector<int> otherShape = other->getShape();
            std::vector<int> newDims = { shape[0], otherShape[1] };

            Tensor* result = new Tensor(newDims, isGrad);

            if(isGrad)
            {
                result->addToPrevious(this);
                result->addToPrevious(other);
                result->gradFn = [result, this, other]() mutable {
                    // Initialize gradients to zero
                std::fill(this->grad.begin(), this->grad.end(), 0.0);
                std::fill(other->grad.begin(), other->grad.end(), 0.0);

                for (int i = 0; i < shape[0]; i++) {
                    for (int j = 0; j < other->getShape()[1]; j++) {
                        for (int k = 0; k < shape[1]; k++) {
                            this->grad[i * strides[0] + k] += result->grad[i * result->strides[0] + j] * other->data[k * other->strides[0] + j];
                            other->grad[k * other->strides[0] + j] += result->grad[i * result->strides[0] + j] * data[i * strides[0] + k];
                        }
                    }
                }
                };
            }
            for(int i =0 ;i< shape[0];i++)
            {
                for(int j=0;j<otherShape[1];j++)
                {
                    float sum = 0;
                    for(int k =0 ;k< shape[1];k++)
                    {
                        sum+= data[i* strides[0] + k] * other->data[k* other->strides[0] + j];
                    }
                    result->data[i* result->strides[0] + j] = sum;
                }
            }
            return result;
        }

    Tensor* Tensor::matmul3D_(Tensor* other) 
        {
            std::vector<int> otherShape = other->getShape();
            std::vector<int> newDims = { shape[0], shape[1], otherShape[2] };

            Tensor* result =  new Tensor(newDims);

            for(int t=0 ; t< shape[0]; t++)
            {
                for(int i =0 ;i< shape[1];i++)
                {
                    for(int j=0;j<otherShape[2];j++)
                    {
                        float sum = 0;
                        for(int k =0 ;k< shape[2];k++)
                        {
                            sum+= data[t * strides[0] + i* strides[1] + k] * other->data[t * strides[0] + k* other->strides[2] + j];
                        }

                        result->data[t * result->strides[0] + i* result->strides[1] + j] = sum;
                    }
                }
            }
            return result;
        }

    Tensor* Tensor::matmul( Tensor* other)  {
        std::vector<int> dims = other->getShape();

        if(other->shape.size()==2)
        return this->matmul2D_(other);
        else
        return this->matmul3D_(other);
    }

    //Transpose
    Tensor* Tensor::transpose() 
    {
        std::vector<int> newDims = {shape.rbegin(), shape.rend()};
        std::vector<int> newStrides = {strides.rbegin(), strides.rend()};
        Tensor* result = new Tensor(data, newDims, newStrides);
        return result;
    }

    //slice
    Tensor* Tensor::slice(int index, int dim)  {
        if (dim >= shape.size()) {
            std::cout << "Dimension out of bounds" << std::endl;
        }
        
        if (index >= shape[dim]) {
            std::cout << "Index out of bounds" << std::endl;
        }

        std::vector<int> newShape = {shape.begin() + dim +1, shape.end()};
        
        Tensor* result =  new Tensor(newShape);
        
        for(int i =0 ; i < strides[dim];i++)
        {
            result->data[i] = data[i + index * strides[dim]];
        }
        
        return result;
    }

    std::vector<int> Tensor::unravelIndex(int index) const {
        std::vector<int> result(shape.size());
        for (int i = shape.size() - 1; i >= 0; i--) {
            result[i] = index % shape[i];
            index /= shape[i];
        }
        return result;
    }
    
    int Tensor::ravelIndex(const std::vector<int>& indices, const std::vector<int>& strides) const {
        int result = 0;
        for (int i = 0; i < indices.size(); i++) {
            result += indices[i] * strides[i];
        }
        return result;
    }

    //backward
    void Tensor::setGrad(bool val)
    {
        this->isGrad = val;
    }

    void Tensor::setGradFn(std::function<void()> fn)
    {
        this->gradFn = fn;
    }

    void Tensor::setGrad(const std::vector<float>& grad)
    {
        this->grad = grad;
    }

    void Tensor::zeroGrad()
    {
        this->grad = std::vector<float>(data.size(), 0.0f);
    }

    std::vector<float>& Tensor::getGrad()
    {
        return this->grad;
    }

    std::vector<int> Tensor::getStrides()
    {
        return this->strides;
    }

    bool Tensor::isGradNode(){
        return this->isGrad;
    }
    //backward

    std::vector<Tensor*> Tensor::getPrevious()
    {
        return this->previous;
    }

   void topologySort(std::vector<Tensor*>& visited, std::vector<Tensor*>& topoList, Tensor* t)
    {
        if (std::find(visited.begin(), visited.end(), t) == visited.end()) 
        {
            visited.push_back(t); 
            for (Tensor* t1 : t->getPrevious())
            {
                topologySort(visited, topoList, t1);
            }
            topoList.push_back(t); 
        }
    }


   void Tensor::backward()
    {
        this->grad[0] = 1.0;
        std::vector<Tensor*> visited;
        std::vector<Tensor*> topoList;
        topologySort(visited, topoList, this);
        for (int i = topoList.size() - 1; i >= 0; i--)
        {
            Tensor* t = topoList[i];
            if (t->gradFn != nullptr)
            {
                t->gradFn();
            }
            //if(!t->isLeafNode() && t->isGrad)
           // {
            //    delete t;
//}
        }
    }
    
    void Tensor::addToPrevious(Tensor* t)
    {
        this->previous.push_back(t);
    }

    //other Ops

    void Tensor::isLeafNode(bool val)
    {
        this->isLeaf = val;
    }

    bool Tensor::isLeafNode(){
        return this->isLeaf;
    }

    Tensor* Tensor::zeros(const std::vector<int>& dims, bool isGrad) {
        Tensor* result = new Tensor(dims, isGrad);
        return result;
    }

    Tensor* Tensor::ones(const std::vector<int>& dims, bool isGrad) {
        Tensor* result = new Tensor(dims, isGrad);
        for (int i = 0; i < result->data.size(); i++) {
            result->data[i] = 1.0f;
        }
        return result;
    }

    void Tensor::setName(std::string name)
    {
        this->name = name;
    }

    std::string Tensor::getName()
    {
        return this->name;
    }

}