43 lines
1005 B
C
43 lines
1005 B
C
|
#pragma once
|
||
|
|
|
||
|
|
#include <cstdlib>
|
||
|
|
#include <vector>
|
||
|
|
|
||
|
|
// You can modify the data structure as you want
|
||
|
|
struct Tensor {
|
||
|
|
// Alloc memory
|
||
|
|
Tensor(std::vector<int> shape_) {
|
||
|
|
ndim = shape_.size();
|
||
|
|
for (int i = 0; i < ndim; i++) { shape[i] = shape_[i]; }
|
||
|
|
int n = num_elem();
|
||
|
|
buf = (float *) malloc(n * sizeof(float));
|
||
|
|
}
|
||
|
|
|
||
|
|
// Alloc memory and copy
|
||
|
|
Tensor(std::vector<int> shape_, float *buf_) {
|
||
|
|
ndim = shape_.size();
|
||
|
|
for (int i = 0; i < ndim; i++) { shape[i] = shape_[i]; }
|
||
|
|
int n = num_elem();
|
||
|
|
buf = (float *) malloc(n * sizeof(float));
|
||
|
|
for (int i = 0; i < n; ++i) { buf[i] = buf_[i]; }
|
||
|
|
}
|
||
|
|
|
||
|
|
~Tensor() {
|
||
|
|
if (buf != nullptr) free(buf);
|
||
|
|
}
|
||
|
|
|
||
|
|
int num_elem() {
|
||
|
|
int sz = 1;
|
||
|
|
for (int i = 0; i < ndim; i++) sz *= shape[i];
|
||
|
|
return sz;
|
||
|
|
}
|
||
|
|
|
||
|
|
// Pointer to data
|
||
|
|
float *buf = nullptr;
|
||
|
|
|
||
|
|
// Shape of tensor, from outermost dimension to innermost dimension.
|
||
|
|
// e.g., {{1.0, -0.5, 2.3}, {4.3, 5.6, -7.8}} => shape = {2, 3}
|
||
|
|
int ndim = 0;
|
||
|
|
int shape[4];
|
||
|
|
};
|