functions.cpp

#include <iostream>
#include <vector>

#include "tensor.h"

namespace tensor {

    //operations
    //matmul
    Tensor* matmul(Tensor* t1, Tensor* t2) {
        if (t1->getShape().size() == 2 && t2->getShape().size() == 2) {
            return t1->matmul(t2);
        }
        else if(t2->getShape().size() == 2) {
            t2 = t2->unsqueeze(0);
        }
       
    
        std::vector<int> dims1 = t1->getShape();
        std::vector<int> dims2 = t2->getShape();

        std::vector<int> batchDims1(dims1.begin(), dims1.end() - 2);
        std::vector<int> batchDims2(dims2.begin(), dims2.end() - 2);
    
        int batchSize = 1;
        for (int dim : batchDims1) {
            batchSize *= dim;
        }

        int innerSize = 1;
        for (int dim : batchDims2) {
            innerSize *= dim;
        }
    
        std::vector<int> outputShape = batchDims1;
        outputShape.push_back(dims1[dims1.size() - 2]);  
        outputShape.push_back(dims2[dims2.size() - 1]); 
    
        std::vector<Tensor*> batchResults;
        
        t1->resize({batchSize, dims1[dims1.size() - 2], dims1[dims1.size() - 1]});
        t2->resize({innerSize, dims2[dims2.size() - 2], dims2[dims2.size() - 1]});

        if(innerSize ==1)
        {
            t2 = t2->broadcast({batchSize, dims2[dims2.size() - 2], dims2[dims2.size() - 1]});
        }

        for (int b = 0; b < batchSize; b++) {
            Tensor* t1Slice = t1->slice(b, 0);
            Tensor* t2Slice = t2->slice(b, 0);
            Tensor* batchResult = t1Slice->matmul(t2Slice);
            batchResults.push_back(batchResult);
        }
        
        Tensor* result = batchResults[0]->unsqueeze(0);
        for (size_t i = 1; i < batchResults.size(); i++) {
            Tensor* batchI = batchResults[i]->unsqueeze(0);
            Tensor* temp = result->concat(batchI, 0);
            result = temp;
        }
        return result;
    }
    
    //Activation Functions

    Tensor* relu( Tensor* t) {
        return t->relu();
    }

    Tensor* sigmoid(Tensor* t) {
  
        return t->sigmoid();
    }

    int Tensor::argmax() {
        int maxIndex = 0;
        float max = data[0];
        for (int i = 1; i < data.size(); i++) {
            if (data[i] > max) {
                max = data[i];
                maxIndex = i;
            }
        }
        return maxIndex;
    }

    Tensor* Tensor::argmax(int dim) {
        std::vector<int> inputShape = this->getShape();
        if (dim < 0) dim += inputShape.size();
        if (dim < 0 || dim >= inputShape.size()) {
            throw std::invalid_argument("Invalid dimension for argmax");
        }
        
        std::vector<int> outputShape = inputShape;
        outputShape.erase(outputShape.begin() + dim);
        
        int dimSize = inputShape[dim];
        int outerSize = 1;
        for (int i = 0; i < dim; i++) {
            outerSize *= inputShape[i];
        }
        
        int innerSize = 1;
        for (int i = dim + 1; i < inputShape.size(); i++) {
            innerSize *= inputShape[i];
        }
        
        std::vector<float> outputData(outerSize * innerSize, 0);
        
        for (int i = 0; i < outerSize; i++) {
            for (int j = 0; j < innerSize; j++) {
                float maxVal = std::numeric_limits<float>::lowest();
                int maxIdx = 0;
                
                for (int k = 0; k < dimSize; k++) {
                    int index = i * dimSize * innerSize + k * innerSize + j;
                    if (this->data[index] > maxVal) {
                        maxVal = this->data[index];
                        maxIdx = k;
                    }
                }
                
                outputData[i * innerSize + j] = static_cast<float>(maxIdx);
            }
        }
        
        Tensor* result = new Tensor(outputData, outputShape);
        return result;
    }

    /*Tensor* tanh(Tensor* t) {
        Tensor* result = new Tensor(2.0f / *(1.0f + (-2.0f * *t)->exp()) - 1.0f);
        return result;
    }*/

    Tensor* softmax( Tensor* t, int dim) {
        Tensor* expT = t->exp();
        Tensor* sumExpT = expT->sum(dim, true);
        return *expT / sumExpT;
    }

    Tensor* mean( Tensor* t, int dim) {
        return *t->sum(dim) / t->getShape()[dim];
    }

    Tensor* Tensor::mean() {    
        Tensor* result =  *this->sum() / this->numel();
        return result;
    }

    //LOSS FUNCTIONS

    Tensor* mse(Tensor* t1,  Tensor* t2) {
        Tensor* result = ((*t1 - t2)->pow(2))->mean();
        return result;
    }

    Tensor* crossEntropy( Tensor* t1,  Tensor* t2) {
        Tensor* temp1 = t1->log();
        Tensor* temp = *t2 * temp1;
        Tensor* result = temp->sum(1);
        return -*result->mean();
    }

    Tensor *binary_cross_entropy(Tensor *t1, Tensor *t2)
    {
        Tensor *temp1 = t1->log();
        Tensor *temp2 = (1 - *t1)->log();
        Tensor *temp = *(*t2 * temp1) + (1 - *t2)->operator*(temp2);
        Tensor *result = -*temp->sum(1)->mean();
        return result;
    }

    bool save(std::unordered_map<std::string, Tensor*> stateDict, std::string path){
        std::ofstream file(path);
        if(file.is_open()){
            for(auto& param : stateDict){
                file << param.first << std::endl;
                file << param.second->getShape().size() << std::endl;
                for(auto& dim : param.second->getShape()){
                    file << dim << " ";
                }
                file << std::endl;
                for(auto& data : param.second->data){
                    file << data << " ";
                }
                file << std::endl;
            }
            file.close();
            return true;
        }
        return false;
    }

    std::unordered_map<std::string, Tensor*> load(std::string path) {
        std::unordered_map<std::string, Tensor*> stateDict;
        std::ifstream file(path);
        
        if (file.is_open()) {
            std::string line;
            while (std::getline(file, line)) {
                std::string paramName = line;
                
                std::getline(file, line);
                int numDims = std::stoi(line);
                
                std::getline(file, line);
                std::stringstream dimStream(line);
                std::vector<int> shape;
                int dim;
                while (dimStream >> dim) {
                    shape.push_back(dim);
                }
                
                std::getline(file, line);
                std::stringstream dataStream(line);
                std::vector<float> data;
                float value;
                while (dataStream >> value) {
                    data.push_back(value);
                }

                Tensor* tensor = Tensor::fromVector(data, shape);
                tensor->isLeafNode(true);
                stateDict[paramName] = tensor;
            }
            file.close();
        }
        
        return stateDict;
    }
}