[TensorDim] Add initializer list ctor

This patch adds a tensordim
initializer list ctor to easily pass as a functional argument

**Self evaluation:**
1. Build test: [X]Passed [ ]Failed [ ]Skipped
2. Run test: [X]Passed [ ]Failed [ ]Skipped

Signed-off-by: Jihoon Lee <jhoon.it.lee@samsung.com>

diff --git a/nntrainer/tensor/tensor_dim.h b/nntrainer/tensor/tensor_dim.h
index 2070152..0d9f3d2 100644 (file)
--- a/nntrainer/tensor/tensor_dim.h
+++ b/nntrainer/tensor/tensor_dim.h
@@ -33,6 +33,28 @@ public:
     feature_len = 0;
   }
 
+  /**
+   * @brief Construct a new Tensor Dim object
+   *
+   * @param dims std::initialize_list
+   *
+   * formats of {w}, {h, w}, {c, h, w}, {b, c, h, w} are accepted
+   */
+  TensorDim(std::initializer_list<unsigned int> dims) : TensorDim() {
+    int shift_size = MAXDIM - dims.size();
+
+    if (shift_size < 0) {
+      throw std::invalid_argument("[TensorDim] max dimension is 4");
+    }
+
+    unsigned int cnt = 0;
+
+    for (auto &i : dims) {
+      setTensorDim(shift_size + cnt, i);
+      cnt += 1;
+    }
+  }
+
   TensorDim(unsigned int b, unsigned int c, unsigned int h, unsigned int w) :
     TensorDim() {
     setTensorDim(0, b);

diff --git a/test/unittest/unittest_nntrainer_tensor.cpp b/test/unittest/unittest_nntrainer_tensor.cpp
index ba3d901..f9a8350 100644 (file)
--- a/test/unittest/unittest_nntrainer_tensor.cpp
+++ b/test/unittest/unittest_nntrainer_tensor.cpp
@@ -16,6 +16,29 @@
 #include <tensor.h>
 #include <tensor_dim.h>
 
+TEST(nntrainer_TensorDim, ctor_initializer_p) {
+  unsigned int b = 3;
+  unsigned int c = 2;
+  unsigned int h = 4;
+  unsigned int w = 5;
+
+  nntrainer::TensorDim t = {w};
+  EXPECT_EQ(nntrainer::TensorDim(1, 1, 1, w), t);
+
+  t = {h, w};
+  EXPECT_EQ(nntrainer::TensorDim(1, 1, h, w), t);
+
+  t = {c, h, w};
+  EXPECT_EQ(nntrainer::TensorDim(1, c, h, w), t);
+
+  t = {b, c, h, w};
+  EXPECT_EQ(nntrainer::TensorDim(b, c, h, w), t);
+}
+
+TEST(nntrainer_TensorDim, ctor_initializer_n) {
+  EXPECT_THROW(nntrainer::TensorDim t({1, 2, 3, 4, 5}), std::invalid_argument);
+}
+
 TEST(nntrainer_TensorDim, setTensorDim_01_p) {
   int status = ML_ERROR_NONE;