tensorops.cpp

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

namespace tensor{

    Tensor* Tensor::concat(Tensor* other, int dim) {

        for (int i = 0; i < shape.size(); ++i) {
            if (i != dim && shape[i] != other->getShape()[i]) {
                throw std::invalid_argument("Incompatible shapes for concatenation");
            }
        }
    
        std::vector<int> newShape = shape;
        newShape[dim] += other->getShape()[dim];

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

        int numElementsBeforeDim = 1;
        for (int i = 0; i < dim; ++i) {
            numElementsBeforeDim *= shape[i];
        }
    
        int numElementsAfterDim = 1;
        for (int i = dim + 1; i < shape.size(); ++i) {
            numElementsAfterDim *= shape[i];
        }
    
        int thisDimSize = shape[dim];
        int otherDimSize = other->getShape()[dim];
    
        // Copy data from this tensor and the other tensor
        for (int i = 0; i < numElementsBeforeDim; ++i) {
            for (int j = 0; j < thisDimSize; ++j) {
                int srcIndex = i * thisDimSize * numElementsAfterDim + j * numElementsAfterDim;
                int destIndex = i * (thisDimSize + otherDimSize) * numElementsAfterDim + j * numElementsAfterDim;
                std::copy(data.begin() + srcIndex, data.begin() + srcIndex + numElementsAfterDim,
                          result->data.begin() + destIndex);
            }
    
            for (int j = 0; j < otherDimSize; ++j) {
                int srcIndex = i * otherDimSize * numElementsAfterDim + j * numElementsAfterDim;
                int destIndex = i * (thisDimSize + otherDimSize) * numElementsAfterDim + (thisDimSize + j) * numElementsAfterDim;
                std::copy(other->data.begin() + srcIndex, other->data.begin() + srcIndex + numElementsAfterDim,
                          result->data.begin() + destIndex);
            }
        }
    
        return result;
    }

    Tensor* Tensor::squeeze(int dim)  {
        if (dim >= shape.size()) {
            std::cout << "Dimension out of bounds" << std::endl;
        }
        if (shape[dim] != 1) {
            std::cout << "Dimension is not 1" << std::endl;
        }

        std::vector<int> newDims = shape;
        newDims.erase(newDims.begin() + dim);
        Tensor* result = new Tensor(data, newDims, isGrad);
        if(isGrad)
        {
            result->addToPrevious(this);
            result->gradFn = [result, this, dim]() mutable {
                for (int i = 0; i < this->data.size(); i++) {
                    this->grad[i] += result->grad[i];
                }
            };
        }

        return result;
    }

    Tensor* Tensor::unsqueeze(int dim)  {
        if (dim > shape.size()) {
            std::cout << "Dimension out of bounds" << std::endl;
        }

        std::vector<int> newDims = shape;
        newDims.insert(newDims.begin() + dim, 1);

        Tensor* result =  new Tensor(data, newDims, isGrad);
        if(isGrad)
        {
            result->addToPrevious(this);
            result->gradFn = [result, this, dim]() mutable {
                for (int i = 0; i < this->data.size(); i++) {
                    this->grad[i] += result->grad[i];
                }
            };
        }
        return result;
    }

    int Tensor::numel() const {
        int result = 1;
        for (int i = 0; i < shape.size(); i++) {
            result *= shape[i];
        }
        return result;
    }
    
    void Tensor::resize(const std::vector<int>& newDims) {
        shape = newDims;
        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]);
    }

    float Tensor::get(const std::vector<int>& indices) const {
        int index = 0;
        for (int i = 0; i < indices.size(); i++) {
            index += indices[i] * strides[i];
        }
        return data[index];
    }

    void Tensor::set(const std::vector<int>& indices, float value) {
        int index = 0;
        for (int i = 0; i < indices.size(); i++) {
            index += indices[i] * strides[i];
        }
        data[index] = value;
    }

    std::vector<int> Tensor::getShape() const {
        return shape;
    }

    float Tensor::at(const std::vector<int>& indices) const {
        if (indices.size() != shape.size()) {
            throw std::runtime_error("Number of indices doesn't match tensor dimensions");
        }
        
        int index = 0;
        for (size_t i = 0; i < indices.size(); i++) {
            if (indices[i] >= shape[i]) {
                throw std::runtime_error("Index out of bounds");
            }
            index += indices[i] * strides[i];
        }
        return data[index];
    }
}