pymoose.argmax#

pymoose.argmax(x, axis, upmost_index, placement=None)[source]#

Compute the index of the maximal element of a tensor along a given dimension.

Parameters

  • x – A tensor.

  • axis – The dimension along which to compute the argmax. upmost_index.

  • upmost_index – The max index that should be considered for computing the argmax. Generally, this should be the size of the axis dimension of x.

  • placement – An optional Placement to pin this operation to.

Returns

A dimension-reduced tensor representing the argmax of x along axis.