Python – tensorflow.gather_nd()

0 0 1 minute read

TensorFlow is open-source Python library designed by Google to develop Machine Learning models and deep learning neural networks.

gather_nd() is used to gather the slice from input tensor based on the indices provided.

Syntax: tensorflow.gather_nd( params, indices, batch_dims, name)

Parameters:

params: It is a Tensor with rank greater than or equal to axis+1.

indices: It is a Tensor of dtype int32 or int64.

batch_dims: It is an integer representing the number o batch dimension. It must be less than rank(indices).

name: It defines the name for the operation.

Returns:

It returns a Tensor having same dtype as param.

Example 1:

Python3

# Importing the library 
import tensorflow as tf 
  
# Initializing the input 
data = tf.constant([[1, 2], [3, 4], [5, 6]]) 
indices = tf.constant([[1], [0], [1]]) 
  
# Printing the input 
print('data: ',data) 
print('indices: ',indices) 
  
# Calculating result 
res = tf.gather_nd(data, indices) 
  
# Printing the result 
print('res: ',res) 

Output:

data:  tf.Tensor(
[[1 2]
 [3 4]
 [5 6]], shape=(3, 2), dtype=int32)
indices:  tf.Tensor(
[[1]
 [0]
 [1]], shape=(3, 1), dtype=int32)
res:  tf.Tensor(
[[3 4]
 [1 2]
 [3 4]], shape=(3, 2), dtype=int32)

Example 2:

Python3

# Importing the library 
import tensorflow as tf 
  
# Initializing the input 
data = tf.constant([[1, 2, 3], [3, 4, 5], [5, 6, 7]]) 
indices = tf.constant([[1, 0], [0, 2], [1, 2]]) 
  
# Printing the input 
print('data: ',data) 
print('indices: ',indices) 
  
# Calculating result 
res = tf.gather_nd(data, indices) 
  
# Printing the result 
print('res: ',res) 

Output:

data:  tf.Tensor(
[[1 2 3]
 [3 4 5]
 [5 6 7]], shape=(3, 3), dtype=int32)
indices:  tf.Tensor(
[[1 0]
 [0 2]
 [1 2]], shape=(3, 2), dtype=int32)
res:  tf.Tensor([3 3 5], shape=(3,), dtype=int32)