argmax() method in C++ Eigen Library

Question

I am using Eigen library for matrix/tensor computation where I want to returns the indices of the maximum values along the depth axis. Similar to what numpy.argmax() does in Python.

Tensor dimension is as follows: (rows = 200, columns = 200, depth=4)

#include <Eigen/Dense>
int main(){
   Eigen::Tensor<double, 3> table(4,200,200);
   table.setRandom();
   // How can I do this task for axis = 2, i.e depth of a tensor?
   // int max_axis = table.argmax(ax=2); 
   return 0;
}

Answer 1

Eigen's tensor library has an argmin/argmax member function, which is unfortunately currently not documented on https://eigen.tuxfamily.org/dox/unsupported/eigen_tensors.html.

Eigen's matrix library can mimic the same behavior via visitor overloads of minCoeff/maxCoeff. See: https://eigen.tuxfamily.org/dox/group__TutorialReductionsVisitorsBroadcasting.html

#include <Eigen/Dense>
#include <unsupported/Eigen/CXX11/Tensor>
#include <iostream>

#define STR_(x)  #x
#define STR(x)  STR_(x)
#define PRINT(x)  std::cout << STR(x) << ":\n" << (x) << std::endl

int main()
{
    using namespace Eigen;
    using T = int;
    using S = Sizes<2, 3>;

    S const sizes{};

    T constexpr data[S::total_size]{
        8, 4,
        1, 6,
        9, 2,
    };

    Map<MatrixX<T> const> const matrix(data, sizes[0], sizes[1]);
    PRINT(matrix);

    RowVector2<Index> argmax{};
    matrix.maxCoeff(&argmax.x(), &argmax.y());
    PRINT(argmax);

    VectorX<Index> argmax0{matrix.cols()};
    for (Index col = 0; col < matrix.cols(); ++col)
        matrix.col(col).maxCoeff(&argmax0[col]);
    PRINT(argmax0);

    VectorX<Index> argmax1{matrix.rows()};
    for (Index row = 0; row < matrix.rows(); ++row)
        matrix.row(row).maxCoeff(&argmax1[row]);
    PRINT(argmax1);

    TensorMap<Tensor<T const, S::count>> const tensor(data, sizes);
    PRINT(tensor);
    PRINT(tensor.argmax());
    PRINT(tensor.argmax(0));
    PRINT(tensor.argmax(1));

    // Note that tensor.argmax() is the index for a 1D view of the data:
    Index const matrix_index = sizes.IndexOfColMajor(std::array{argmax.x(), argmax.y()});
    Index const tensor_index = Tensor<Index, 0>{tensor.argmax()}();
    PRINT(matrix_index == tensor_index);
}

Outputs:

matrix:
8 1 9
4 6 2
argmax:
0 2
argmax0:
0
1
0
argmax1:
2
1
tensor:
8 1 9
4 6 2
tensor.argmax():
4
tensor.argmax(0):
0
1
0
tensor.argmax(1):
2
1
matrix_index == tensor_index:
1