constructor.hpp

#pragma once
 
#include "lamppp/tensor/scalar.hpp"  // TODO : maybe move scalar somewhere ?
#include "variable.hpp"
 
namespace lmp::autograd {
 
using std::multiplies;
 
Variable zeros(const std::vector<size_t>& shape, tensor::DeviceType device,
               tensor::DataType dtype, bool requires_grad);
 
Variable ones(const std::vector<size_t>& shape, tensor::DeviceType device,
              tensor::DataType dtype, bool requires_grad);
 
Variable rand(const std::vector<size_t>& shape, tensor::DeviceType device,
              tensor::DataType dtype, bool requires_grad);
 
Variable randn(tensor::Scalar mean, tensor::Scalar var, const std::vector<size_t>& shape, tensor::DeviceType device,
              tensor::DataType dtype, bool requires_grad);
 
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, tensor::DeviceType device,
                tensor::DataType dtype, bool requires_grad) {
  TensorHelper constr;
  constr.unroll(data);
  return Variable(tensor::Tensor(constr.data, constr.shape, device, dtype),
                  requires_grad);
}
 
}  // namespace lmp::autograd