Implement top_k for Eigen in c++

Question

I have a 3D Eigen tensor (Eigen::Tensor<float, 3, Eigen::RowMajor, EIGEN_DEFAULT_DENSE_INDEX_TYPE>) of shape [B, D, C] containing floating point numbers that represent the confidence score of some network output.

I want to perform top k and obtain the values and indices of the highest-scoring outputs along the second dimension and reduce D to a smaller dimension k. The resulting vector should be of shape [B, k, C].

In python (tensorflow) it would simply be (in this case integers for simplicity):

data = np.array([[[3, 9, 8, 5],
        [3, 3, 0, 1],
        [5, 5, 1, 9]],

       [[1, 3, 0, 9],
        [8, 3, 7, 7],
        [8, 0, 9, 5]]])


values, indices = tf.math.top_k(tensor, k=2, sorted=False)

And the result:

Values

[[[8 9]
  [3 3]
  [5 9]]

 [[3 9]
  [7 8]
  [8 9]]]

Indices

[[[2 1]
  [1 0]
  [0 3]]

 [[1 3]
  [2 0]
  [0 2]]]

Thanks in advance!

Answer 1

The following code works, but it is not efficient:

#include <Eigen/Dense>
#include <algorithm> 
#include <vector>

int B = tensor.dimension(0);
int D = tensor.dimension(1);
int C = tensor.dimension(2);
int k = 2; // Example value for top-k

// Resulting tensors
Eigen::Tensor<float, 3, Eigen::RowMajor> values(B, k, C);
Eigen::Tensor<int, 3, Eigen::RowMajor> indices(B, k, C);

for (int b = 0; b < B; ++b) {
    for (int c = 0; c < C; ++c) {
        // Extract a slice for the current batch element and class
        Eigen::Tensor<float, 1, Eigen::RowMajor> slice = data.chip(b, 0).chip(c, 1);

        // Vector to store index-value pairs 
        std::vector<std::pair<float, int>> index_value_pairs;
        index_value_pairs.reserve(D);  // Pre-allocation for efficiency

        // Store index-value pairs
        for (int d = 0; d < D; ++d) {
            index_value_pairs.emplace_back(slice(d), d);
        }

        // Partial sort (only the top-k) based on the values
        std::partial_sort(index_value_pairs.begin(), index_value_pairs.begin() + k, index_value_pairs.end(),
                          [](const std::pair<float, int>& a, const std::pair<float, int>& b) {
                              return a.first > b.first; // Sort by value in descending order
                          });

        // Fill the results 
        for (int i = 0; i < k; ++i) {
            values(b, i, c) = index_value_pairs[i].first;
            indices(b, i, c) = index_value_pairs[i].second;
        }

    }
}