Tensorflow.js tf.gather() Function

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

The .gather() function is used to collect the fragments from the stated tensor x’s axis as per the stated indices.

Syntax:  

tf.gather(x, indices, axis?, batchDims?)

Parameters:  

• x: It is the stated input tensor whose fragments are to be collected, and it can be of type tf.Tensor, TypedArray, or Array.
• indices: It is the stated indices of the values which are to be pulled out, and it can be of type tf.Tensor, TypedArray, or Array.
• axis: It is the stated axis above which the values are to be selected. The by default value is zero, and it is of type number. However, this parameter is optional.
• batchDims: It is the stated number of batch sizes and should be less than or equal to stated rank i.e. indices. It’s by default value is zero. Moreover, the output returned must have the shape of x.shape[:axis] + indices.shape[batchDims:] + x.shape[axis + 1:]. It is of type number and is optional.

Return Value: It returns tf.Tensor object.

Example 1:

Javascript

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

// Defining tensor input and indices
const y = tf.tensor1d([1, 6, 7, 8]);
const ind = tf.tensor1d([1, 6, 2], 'int32');

// Calling tf.gather() method and
// Printing output
y.gather(ind).print();

Output:

Tensor
    [6, NaN, 7]

Example 2:

Javascript

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

// Defining tensor input, indices, axis,
// and batchdims
const y = tf.tensor2d([7, 8, 12, 13], [4, 1]);
const ind = tf.tensor1d([2, 3, 0], 'int32');
const axis = 1;
const batchdims = -1;

// Calling tf.gather() method
var res = tf.gather(y, ind, axis, batchdims);

// Printing output
res.print();

Output:

Tensor
    [[12 , 13 , 7 ],
     [13 , NaN, 8 ],
     [NaN, NaN, 12],
     [NaN, NaN, 13]]

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