oop_and_ml/src/tensor.h

59 lines
No EOL
1.8 KiB
C++
Executable file

#ifndef TENSOR_H
#define TENSOR_H
#include <stdio.h>
#include <vector>
#include <iostream>
class tensor {
public:
tensor(); // scalar 0
explicit tensor(double v); // scalar v
tensor(int dim, size_t shape[], double data[]); // From C
tensor(int dim, std::initializer_list<size_t> shape, std::initializer_list<double> data); // Conv2d
int get_dim() const;
// scalar only
double item() const;
double &item();
// more accessors
double at(size_t i) const;
double at(size_t i, size_t j) const;
double at(size_t i, size_t j, size_t k, size_t l) const;
double& at(size_t i, size_t j, size_t k, size_t l);
// passing tensor back to C code
size_t *get_shape_array();
double *get_data_array();
// easy way to print out the tensor for logging
void print() const;
size_t getTotalSize() const;
// friend is better than member, because it allows
// double + tensor (implicit conversion)
// operator overloads
friend tensor operator+(const tensor &lhs, const tensor &rhs);
friend tensor operator-(const tensor &lhs, const tensor &rhs);
friend tensor operator*(const tensor &lhs, const tensor &rhs);
template <class T>
friend tensor helper(const tensor &lhs, const tensor &rhs);
friend tensor ReLU(const tensor &x);
friend tensor Flatten(const tensor &x);
friend tensor Input2d(const tensor &x);
friend tensor Linear(const tensor &w, const tensor &x, const tensor &b);
friend tensor MaxPool2d(const tensor &x, const int kernel_size, const int stride);
friend tensor Conv2D(const tensor &input_tensor, const tensor &kernel_tensor, const tensor &bias_tensor, int in_channels, int out_channels, int kernel_size);
private:
int dim_;
std::vector<size_t> shape_;
std::vector<double> data_;
}; // class tensor
#endif // TENSOR_H;