Function argMax

Returns the indices of the maximum values along an axis.

The result has the same shape as input with the dimension along axis removed.

const x = tf.tensor1d([1, 2, 3]);

x.argMax().print();  // or tf.argMax(x)
Copy

const x = tf.tensor2d([1, 2, 4, 3], [2, 2]);

const axis = 1;
x.argMax(axis).print();  // or tf.argMax(x, axis)
Copy

Doc