Tensor.h

#pragma once
#include <vector>
#include <string>
#include <unordered_map>
#include <iostream>
#include <memory>
#include <queue>
#include <stdexcept>
#include "optypes.h"

enum class DType {
  FLOAT,
  INT32,
  INT64
};

template <typename T>
class Tensor {
public:
  Tensor() : dtype_(DType::FLOAT) {}

  Tensor(std::vector<size_t> shape, DType dtype)
      : shape_(std::move(shape)), dtype_(dtype), data_(size()) {}

  Tensor(std::vector<size_t> shape, DType dtype, std::vector<T> data)
      : shape_(std::move(shape)), dtype_(dtype), data_(std::move(data)) {
    if (data_.size() != size()) {
      throw std::runtime_error("Tensor: data size does not match shape");
    }
  }

  Tensor(std::vector<size_t> shape, DType dtype, const T* ptr)
      : shape_(std::move(shape)), dtype_(dtype), data_(size()) {
    std::copy(ptr, ptr + data_.size(), data_.begin());
  }

  const std::vector<size_t>& shape() const { return shape_; }
  DType dtype() const { return dtype_; }

  T* data() { return data_.data(); }
  const T* data() const { return data_.data(); }

  size_t size() const {
    size_t s = 1;
    for (size_t dim : shape_) s *= dim;
    return s;
  }

  std::vector<T>& vec() { return data_; }
  const std::vector<T>& vec() const { return data_; }

private:
  std::vector<size_t> shape_;
  std::vector<T> data_;
  DType dtype_;
};