Tensorflow.js tf.argMax() Function

Tensorflow.js is an open-source library developed by Google for running machine learning models and deep learning neural networks in the browser or node environment.

The tf.argMax() function is used to return the indices for the maximum values of the specified Tensor along an axis.

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

Syntax:

tf.argMax (x, axis)

Parameters: This function accepts two parameters which are illustrated below:

• x: The input tensor.
• axis: The specified dimension(s) to reduce. It is an optional parameter and its default value is 0.

Return Value: It returns a Tensor of the indices of the maximum values along an axis.

Example 1:

Javascript

// Importing the tensorflow.js library
import * as tf from "@tensorflow/tfjs"

// Initializing a some tensors
const a = tf.tensor1d([1, 0]);
const b = tf.tensor1d([3, 5]);
const c = tf.tensor1d([6, 3, 5, 12]);

// Calling the .argMax() function over
// the above tensors
a.argMax().print();
b.argMax().print();
c.argMax().print();

Output:

Tensor
   0
Tensor
   1
Tensor
   3

Example 2:

Javascript

// Importing the tensorflow.js library
import * as tf from "@tensorflow/tfjs"

// Initializing a some tensors
const a = tf.tensor1d([0, 1]);
const b = tf.tensor2d([9, 5, 2, 8], [2, 2]);
const c = tf.tensor1d([6, 4, 7]);

// Initializing a axis parameters
const axis1 = -1;
const axis2 = -2;
const axis3 = 0;

// Calling the .argMax() function over
// the above tensors
a.argMax(axis1).print();
b.argMax(axis2).print();
c.argMax(axis3).print();

Output:

Tensor
   1
Tensor
   [0, 1]
Tensor
   2

Reference: https://js.tensorflow.org/api/latest/#argMax