tensor.hpp

#pragma once
 
#include <iostream>
#include <memory>
#include <span>
#include <utility>
#include <vector>
#include "data_type.hpp"
#include "device_type.hpp"
#include "dispatch_type.hpp"
#include "fill_like.hpp"
#include "tensor_impl.hpp"
#include "lamppp/tensor/native/memory_ops.hpp"
 
namespace lmp::tensor {
 
namespace detail {
class UnsafeTensorAccessor;
}
 
class Tensor {
 public:
  Tensor() = default;
 
  template <typename T>
  explicit Tensor(const std::vector<T>& data, const std::vector<size_t>& shape,
                  DeviceType device = DeviceType::CPU,
                  DataType dtype = DataType::Float64)
      : impl_(std::make_shared<TensorImpl>(data, shape, device, dtype)) {}
 
  void* data() const noexcept;
  DataType type() const noexcept;
  DeviceType device() const noexcept;
  const std::vector<size_t>& shape() const noexcept;
  const std::vector<detail::stride_t>& strides() const noexcept;
  size_t numel() const noexcept;
 
  template <typename T>
  std::vector<T> to_vector() const {
    std::vector<T> converted_data(impl_->numel());
    LMP_DISPATCH_ALL_TYPES(impl_->type(), [&] {
      std::unique_ptr<scalar_t[]> original_data = std::make_unique<scalar_t[]>(numel());
      ops::copy_stub()(device(), DeviceType::CPU, data(),
                                  original_data.get(), numel(), type(), type());
 
      for (size_t i = 0; i < impl_->numel(); ++i) {
        converted_data[i] = static_cast<T>(original_data[i]);
      }
    });
    return converted_data;
  }
  Scalar index(const std::vector<size_t>& idx) const;
 
  Tensor reshape(std::vector<size_t> new_shape) const;
  Tensor squeeze(size_t dim) const;
  Tensor expand_dims(size_t dim) const;
 
  Tensor to(DeviceType device) const;
 
  void copy(const Tensor& other);
  void fill(Scalar item);
 
  friend std::ostream& operator<<(std::ostream& os, const Tensor& obj);
  friend class TensorOpFact;
  friend class detail::UnsafeTensorAccessor;
 
 private:
  explicit Tensor(std::shared_ptr<TensorImpl> ptr) : impl_(std::move(ptr)) {}
  std::shared_ptr<TensorImpl> impl_;
};
 
namespace detail {
// @internal
struct UnsafeTensorAccessor {
  static std::shared_ptr<TensorImpl> getImpl(const Tensor& ten) {
    return ten.impl_;
  }
  static Tensor fromImpl(std::shared_ptr<TensorImpl> ptr) {
    return Tensor(std::move(ptr));
  }
};
// @endinternal
}  // namespace detail
 
}  // namespace lmp::tensor