[FIXED] Argmax of 2d vector on C++

Issue

I am working on python/pytorch and I have an example like

2d vector a
 
    |
    v
  dim-0  ---> -----> dim-1 ------> -----> --------> dim-1
    |   [[-1.7739,  0.8073,  0.0472, -0.4084],
    v    [ 0.6378,  0.6575, -1.2970, -0.0625],
    |    [ 1.7970, -1.3463,  0.9011, -0.8704],
    v    [ 1.5639,  0.7123,  0.0385,  1.8410]]
    |
    v

Then, the argmax with the index of 1 will be

# argmax (indices where max values are present) along dimension-1
In [215]: torch.argmax(a, dim=1)
Out[215]: tensor([1, 1, 0, 3])

My question is that given the 2d vector a as above, how could I implement argmax function on C++ to give me same output as above? Thanks for reading This is what I did

   

    vector<vector<float>> a_vect

    {
      {-1.7739,  0.8073,  0.0472, -0.4084},
      {0.6378,  0.6575, -1.2970, -0.0625},
      {1.7970, -1.3463,  0.9011, -0.8704},
      {1.5639,  0.7123,  0.0385,  1.8410}
     };
     std::vector<int>::iterator max = max_element(a_vect.begin() , a_vect.end()-a_vect.begin());

Solution

You can use std::max_element to find the index in each sub vector

#include <algorithm>
#include <iostream>
#include <vector>
using std::vector;

int main()
{
  vector<vector<float>> a_vect=
  {
    {-1.7739,  0.8073,  0.0472, -0.4084},
    {0.6378,  0.6575, -1.2970, -0.0625},
    {1.7970, -1.3463,  0.9011, -0.8704},
    {1.5639,  0.7123,  0.0385,  1.8410}
  };
  
  vector<int> max_index;
  for(auto& v:a_vect)
    max_index.push_back(std::max_element(v.begin(),v.end())-v.begin());
  
  for(auto i:max_index)
    std::cout << i << ' '; // 1 1 0 3 
}

Answered By - apple apple