constructor.hpp

#pragma once
 
#include "lamppp/tensor/data_type.hpp"
#include "lamppp/common/config.hpp"
#include "lamppp/autograd/variable.hpp"
 
namespace lmp::autograd {
 
using std::multiplies;
 
Variable zeros(const std::vector<size_t>& shape, bool requires_grad = false,
               tensor::DeviceType device = DEFAULT_DEVICE,
               tensor::DataType dtype = DEFAULT_DTYPE);
 
Variable ones(const std::vector<size_t>& shape, bool requires_grad = false, tensor::DeviceType device = DEFAULT_DEVICE,
              tensor::DataType dtype = DEFAULT_DTYPE);
 
Variable rand(const std::vector<size_t>& shape, bool requires_grad = false,
              tensor::DeviceType device = DEFAULT_DEVICE,
              tensor::DataType dtype = DEFAULT_DTYPE);
 
Variable randn(tensor::Scalar mean, tensor::Scalar var,
               const std::vector<size_t>& shape, bool requires_grad = false,
               tensor::DeviceType device = DEFAULT_DEVICE,
               tensor::DataType dtype = DEFAULT_DTYPE);
 
template <typename>
struct IsVector : std::false_type {};
template <typename U, typename Alloc>
struct IsVector<std::vector<U, Alloc>> : std::true_type {};
 
struct TensorHelper {
  std::vector<tensor::Scalar> data;
  std::vector<size_t> shape;
  template <typename T>
  void unroll(const std::vector<T>& tensor, size_t depth = 0) {
    if (depth >= shape.size()) {
      shape.push_back(tensor.size());
    }
    LMP_CHECK(tensor.size() == shape[depth])
        << "Dimensions along axis must be consistent.";
    if constexpr (IsVector<T>::value) {
      for (const T& t : tensor) {
        unroll(t, depth + 1);
      }
    } else {
      data.insert(data.end(), tensor.begin(), tensor.end());
    }
  }
};
 
 
template <typename V>
Variable tensor(const std::vector<V>& data, bool requires_grad = false,
                tensor::DeviceType device = DEFAULT_DEVICE,
                tensor::DataType dtype = DEFAULT_DTYPE) {
  TensorHelper constr;
  constr.unroll(data);
  return Variable(tensor::Tensor(constr.data, constr.shape, device, dtype),
                  requires_grad);
}
 
}  // namespace lmp::autograd