[FEA] find_idx with dim

[osayamenja]

Is your feature request related to a problem? Please describe.
find_idx that accepts dim argument.

Describe the solution you'd like

• For an MxN matrix, I would like to call find_idx across the rows.
• Output tensor would now be 2D here, with each row containing the indices.
• num_found would now be a 1D vector with each entry containing the number of matches per row.

Describe alternatives you've considered
I call the current find_idx on each row of my matrix in a loop.

Additional context
We can simplify to support only matrices. Thanks!

[osayamenja]

This may already be supported currently, if so, do let me know!

[cliffburdick]

Hi @osayamenja , thanks for filing this. find_idx currently uses CUB's DeviceSelect for its implementation, and DeviceSelect doesn't have a segmented/batching API like many of their other functions do. Ideally we would be able to batch this so you don't have to call it in a loop, but that's not possible yet.

We will look into adding an index argument so you can do this in a loop, but in the meantime you can accomplish it using slice:

auto t = make_tensor<float>({M, N});
auto idx_array = make_tensor<float>({M});
auto num_found = make_tensor<int>();
for (int i = 0; i < t.Size(0); i++) {
  auto s = slice(t, {i, 0}, {matxDropDim, matxEnd}); // Slice row i
  (mtie(idx_array, num_found) = find_idx(s, GT{0.5f})).run(exec);
  // Do something with `num_found` and `-idx_array`
}
  

[osayamenja]

@cliffburdick Thanks so much for the detailed and quick response, will close!