tensor.h

#ifndef TENSOR_H
#define TENSOR_H
 
#include <cassert>
#include <iostream>
#include <memory>
#include <string>
#include <vector>
 
enum class Backend { CPU, CUDA };
 
class Tensor {
public:
    class Impl;
    class CPUImpl;
    class CUDAImpl;
 
    class View;
    class Shape;
 
public:
    Tensor(const Tensor::Shape& shape, Backend backend = Backend::CPU);
 
    Tensor(const std::initializer_list<size_t>& shape, Backend backend = Backend::CPU);
 
    Tensor(const Tensor::Shape& shape, float value, Backend backend = Backend::CPU);
 
    Tensor(const std::initializer_list<size_t>& shape, float value, Backend backend = Backend::CPU);
 
    Tensor(float value, Backend backend = Backend::CPU);
 
    Tensor(const Tensor& tensor);
 
    Tensor(std::unique_ptr<Tensor::Impl> impl, Backend backend = Backend::CPU);
 
    ~Tensor();
 
    void to(Backend newBackend);
 
    void fillAll(float value);
 
    void fillRand();
 
    void print() const;
 
    void print(const std::vector<size_t>& position) const;
 
    Tensor::Shape getShape() const;
 
    std::string getBackendString() const;
 
    inline Backend getBackend() const { return m_backend; }
 
    std::string getDataString() const;
 
    size_t getNumElements() const;
 
    std::vector<float> toVector() const;
 
    void set(const std::vector<size_t>& position, const Tensor::View& rhs);
 
    bool compare(const Tensor::View& rhs) const;
 
    bool compare(const std::vector<size_t>& position, const Tensor::View& rhs) const;
 
    Tensor operator+(const Tensor::View& rhs) const;
 
    Tensor operator-(const Tensor::View& rhs) const;
 
    Tensor operator*(const Tensor::View& rhs) const;
 
    Tensor operator/(const Tensor::View& rhs) const;
 
    Tensor operator+(float scalar) const;
 
    Tensor operator-(float scalar) const;
 
    Tensor operator*(float scalar) const;
 
    Tensor operator/(float scalar) const;
 
    friend Tensor operator+(float scalar, const Tensor& rhs);
 
    friend Tensor operator-(float scalar, const Tensor& rhs);
 
    friend Tensor operator*(float scalar, const Tensor& rhs);
 
    friend Tensor operator/(float scalar, const Tensor& rhs);
 
    void operator+=(const Tensor::View& rhs);
 
    void operator-=(const Tensor::View& rhs);
 
    void operator*=(const Tensor::View& rhs);
 
    void operator/=(const Tensor::View& rhs);
 
    Tensor mean(size_t dim = 0) const;
 
    Tensor sum(size_t dim = 0) const;
 
    Tensor min(size_t dim = 0) const;
 
    Tensor max(size_t dim = 0) const;
 
    Tensor prod(size_t dim = 0) const;
 
    Tensor norm() const;
 
    Tensor all(size_t dim = 0) const;
 
    Tensor any(size_t dim = 0) const;
 
    Tensor matmul(const Tensor::View& rhs) const;
 
    Tensor::View operator[](size_t idx) const;
 
    Tensor::View subsample(std::vector<size_t> strides) const;
 
    Tensor& operator=(const Tensor& rhs);
 
    Tensor& operator=(const Tensor::View& rhs);
 
    Tensor& operator=(float scalar);
 
    bool isEqual(const Tensor::View& rhs) const;
 
    bool isNotEqual(const Tensor::View& rhs) const;
 
    Tensor operator==(const Tensor::View& rhs) const;
 
    Tensor operator!=(const Tensor::View& rhs) const;
 
    Tensor operator<(const Tensor::View& rhs) const;
 
    Tensor operator<=(const Tensor::View& rhs) const;
 
    Tensor operator>(const Tensor::View& rhs) const;
 
    Tensor operator>=(const Tensor::View& rhs) const;
 
    Tensor isClose(const Tensor::View& rhs, float tolerance = 1e-5f) const;
 
private:
    Backend m_backend;
    std::unique_ptr<Impl> m_impl;
 
    template <typename BinaryOp>
    Tensor applyBinaryOp(const Tensor::View& rhs, BinaryOp op) const;
 
    template <typename BinaryOp>
    void applyInplaceBinaryOp(const Tensor::View& rhs, BinaryOp op);
 
    template <typename ReductionOp>
    Tensor applyReduction(size_t dim, ReductionOp op) const;
};
 
#endif  // TENSOR_H