TensorFlow-style Tensor

1#include <vector>
2#include <memory>
3#include <numeric>
4 
5enum class DataType {
6  FLOAT32,
7  FLOAT64,
8  INT32,
9  INT64,
10  BOOL
11};
12 
13class TensorShape {
14  std::vector<int64_t> dims_;
15 
16public:
17  TensorShape() = default;
18  TensorShape(std::initializer_list<int64_t> dims) : dims_(dims) {}
19 
20  int num_dims() const { return static_cast<int>(dims_.size()); }
21  int64_t dim_size(int d) const { return dims_[d]; }
22 
23  int64_t num_elements() const {
24    if (dims_.empty()) return 0;
25    return std::accumulate(dims_.begin(), dims_.end(), 
26                           int64_t{1}, std::multiplies<int64_t>());
27  }
28};
29 
30class Tensor {
31  DataType dtype_;
32  TensorShape shape_;
33  std::shared_ptr<void> buffer_;
34 
35public:
36  Tensor(DataType dtype, const TensorShape& shape)
37    : dtype_(dtype), shape_(shape) {}
38 
39  DataType dtype() const { return dtype_; }
40  const TensorShape& shape() const { return shape_; }
41  int dims() const { return shape_.num_dims(); }
42  int64_t NumElements() const { return shape_.num_elements(); }
43};