def Flat.argmax_rows

bend-tensors@0.0.0.2 · bend_tensors.bend:2095

import Base

2095# each row's argmax, in row order; the matrix is consumed
2096def Flat.argmax_rows(+rows: Nat, +cols: Nat, m: Flat<rows, cols>) -> List<&2, U32>:
2097  FlatS{a} = m
2098  Flat.argmax.of(rows, cols, a)