Función Tensorflow.js tf.gather()

Tensorflow.js es una biblioteca de código abierto desarrollada por Google para ejecutar modelos de aprendizaje automático, así como redes neuronales de aprendizaje profundo en el entorno del navegador o del Node.

La función .gather() se usa para recolectar los fragmentos del eje del tensor x establecido según los índices establecidos.

Sintaxis:  

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

Parámetros:  

• x: Es el tensor de entrada indicado cuyos fragmentos se van a recolectar, y puede ser de tipo tf.Tensor, TypedArray o Array.
• índices: son los índices establecidos de los valores que se extraerán, y puede ser del tipo tf.Tensor, TypedArray o Array.
• eje: Es el eje indicado sobre el cual se seleccionarán los valores. El valor por defecto es cero, y es de tipo número. Sin embargo, este parámetro es opcional.
• batchDims: es el número indicado de tamaños de lote y debe ser menor o igual al rango indicado, es decir, índices. Su valor predeterminado es cero. Además, la salida devuelta debe tener la forma de x.shape[:axis] + indices.shape[batchDims:] + x.shape[axis + 1:]. Es de tipo número y es opcional.

Valor devuelto: Devuelve el objeto tf.Tensor.

Ejemplo 1:

// 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();

Producción:

Tensor
    [6, NaN, 7]

Ejemplo 2:

// 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();

Producción:

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

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