mirror of
https://github.com/wassname/ray.git
synced 2026-07-09 22:39:18 +08:00
Philipp Moritz
51 lines
1.7 KiB
C++
51 lines
1.7 KiB
C++
#include "tensor.h"
|
|
|
|
using namespace arrow;
|
|
|
|
namespace numbuf {
|
|
|
|
template<typename T>
|
|
TensorBuilder<T>::TensorBuilder(const TypePtr& dtype, MemoryPool* pool)
|
|
: dtype_(dtype) {
|
|
dim_data_ = std::make_shared<Int64Builder>(pool);
|
|
dims_ = std::make_shared<ListBuilder>(pool, dim_data_);
|
|
value_data_ = std::make_shared<PrimitiveBuilder<T>>(pool, dtype);
|
|
values_ = std::make_shared<ListBuilder>(pool, value_data_);
|
|
auto dims_field = std::make_shared<Field>("dims", dims_->type());
|
|
auto values_field = std::make_shared<Field>("data", values_->type());
|
|
auto type = std::make_shared<StructType>(std::vector<FieldPtr>({dims_field, values_field}));
|
|
tensors_ = std::make_shared<StructBuilder>(pool, type, std::vector<std::shared_ptr<ArrayBuilder>>({dims_, values_}));
|
|
};
|
|
|
|
template<typename T>
|
|
Status TensorBuilder<T>::Append(const std::vector<int64_t>& dims, const elem_type* data) {
|
|
RETURN_NOT_OK(tensors_->Append());
|
|
RETURN_NOT_OK(dims_->Append());
|
|
RETURN_NOT_OK(values_->Append());
|
|
int32_t size = 1;
|
|
for (auto dim : dims) {
|
|
size *= dim;
|
|
RETURN_NOT_OK(dim_data_->Append(dim));
|
|
}
|
|
RETURN_NOT_OK(value_data_->Append(data, size));
|
|
return Status::OK(); // tensors_->Append();
|
|
}
|
|
|
|
template<typename T>
|
|
std::shared_ptr<Array> TensorBuilder<T>::Finish() {
|
|
return tensors_->Finish();
|
|
}
|
|
|
|
template class TensorBuilder<UInt8Type>;
|
|
template class TensorBuilder<Int8Type>;
|
|
template class TensorBuilder<UInt16Type>;
|
|
template class TensorBuilder<Int16Type>;
|
|
template class TensorBuilder<UInt32Type>;
|
|
template class TensorBuilder<Int32Type>;
|
|
template class TensorBuilder<UInt64Type>;
|
|
template class TensorBuilder<Int64Type>;
|
|
template class TensorBuilder<FloatType>;
|
|
template class TensorBuilder<DoubleType>;
|
|
|
|
}
|