Tensorflow.js tf.topk() 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.topk() function along with the last dimension is also used in finding the values and indices of the k largest entries.

Syntax:

tf.topk (x, k?, sorted?)

Parameters:

• x: 1-D or higher tf.Tensor with last dimension being at least k.
• k: It is the number of elements to look for.
• sorted: It is the boolean value. If it is true, then the resulting k elements will be sorted by the values in descending order.

Return Value: {values: tf.Tensor, indices: tf.Tensor}. It returns an object with two tensors containing the values and indices.

Example 1:

Javascript

const tf = require("@tensorflow/tfjs")

// Creating a 2d tensor
const a = tf.tensor2d([[1, 20, 3], [4, 3, 1], [8, 9, 10]]);
const {values, indices} = tf.topk(a);

// Printing the values and indices
values.print();
indices.print();

Output:

Tensor
    [[20],
     [4 ],
     [10]]
Tensor
    [[1],
     [0],
     [2]]

Example 2: In this example, we will provide the argument k, to get the largest k entries.

Javascript

const tf = require("@tensorflow/tfjs")

// Creating a 2d tensor
const a = tf.tensor2d([[1, 20, 3], [4, 3, 1], [8, 9, 10]]);
const {values, indices} = tf.topk(a, 3);

// Printing the values and indices
values.print();
indices.print();

Output:

As we have passed k = 3, we get 3 largest values in our result.

Tensor
    [[20, 3, 1],
     [4 , 3, 1],
     [10, 9, 8]]
Tensor
    [[1, 2, 0],
     [0, 1, 2],
     [2, 1, 0]]

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