runtime/libs/ndarray/example/example_array.cpp

   1 /*
   2  * Copyright (c) 2019 Samsung Electronics Co., Ltd. All Rights Reserved
   3  *
   4  * Licensed under the Apache License, Version 2.0 (the "License");
   5  * you may not use this file except in compliance with the License.
   6  * You may obtain a copy of the License at
   7  *
   8  *      http://www.apache.org/licenses/LICENSE-2.0
   9  *
  10  * Unless required by applicable law or agreed to in writing, software
  11  * distributed under the License is distributed on an "AS IS" BASIS,
  12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13  * See the License for the specific language governing permissions and
  14  * limitations under the License.
  15  */
  16
  17 #include "ndarray/Array.h"
  18
  19 #include <iostream>
  20 #include <iterator>
  21
  22 using namespace ndarray;
  23
  24 void gather_array(const Array<float> &input, Array<float> &output, const Array<int> &indices)
  25 {
  26   assert(indices.shape().rank() == 3);
  27   assert(input.shape().rank() == 3);
  28   assert(indices.shape().dim(1) == input.shape().rank());
  29
  30   for (size_t i = 0; i < indices.shape().dim(0); ++i)
  31   {
  32     for (size_t j = 0; j < indices.shape().dim(1); ++j)
  33     {
  34       auto index = indices.slice(i, j);
  35       output.slice(i, j).assign(input.slice(index[0], index[1]));
  36     }
  37   }
  38 }
  39
  40 int main()
  41 {
  42   // fill tensor of shape[3,3,4] with sequential numbers from [0..36)
  43   Shape in_shape{3, 3, 4};
  44   std::vector<float> input_data(in_shape.element_count());
  45   for (size_t i = 0; i < in_shape.element_count(); ++i)
  46     input_data[i] = i;
  47
  48   Array<float> input(input_data.data(), in_shape);
  49
  50   // select column-vectors on main diagonal
  51   Shape indices_shape{1, 3, 2};
  52   std::vector<int> indices_data(indices_shape.element_count());
  53   Array<int> indices(indices_data.data(), indices_shape);
  54
  55   indices.slice(0, 0) = {0, 0};
  56   indices.slice(0, 1) = {1, 1};
  57   indices.slice(0, 2) = {2, 2};
  58
  59   Shape output_shape{1, 3, 4};
  60   std::vector<float> output_data(output_shape.element_count());
  61
  62   Array<float> output(output_data.data(), output_shape);
  63
  64   gather_array(input, output, indices);
  65
  66   for (size_t i = 0; i < indices_shape.dim(0); ++i)
  67   {
  68     for (size_t j = 0; j < indices_shape.dim(1); ++j)
  69     {
  70       auto output_piece = output.slice(i, j);
  71       std::ostream_iterator<int> cout_it(std::cout, ", ");
  72       std::copy(output_piece.begin(), output_piece.end(), cout_it);
  73       std::cout << std::endl;
  74     }
  75   }
  76 }