mon
mon

Reputation: 22274

Tensorflow 2 - what is 'index depth' in tensor_scatter_nd_update?

Please explain what is index depth of tf.tensor_scatter_nd_update.

tf.tensor_scatter_nd_update(
    tensor, indices, updates, name=None
)

Why indices is 2D for 1D tensor?

indices has at least two axes, the last axis is the depth of the index vectors. For a higher rank input tensor scalar updates can be inserted by using an index_depth that matches tf.rank(tensor):

tensor = [0, 0, 0, 0, 0, 0, 0, 0]    # tf.rank(tensor) == 1
indices = [[1], [3], [4], [7]]       # num_updates == 4, index_depth == 1   # <--- what is depth and why 2D for 1D tensor?
updates = [9, 10, 11, 12]            # num_updates == 4
print(tf.tensor_scatter_nd_update(tensor, indices, updates))

enter image description here

tensor = [[1, 1], [1, 1], [1, 1]]    # tf.rank(tensor) == 2
indices = [[0, 1], [2, 0]]           # num_updates == 2, index_depth == 2
updates = [5, 10]                    # num_updates == 2
print(tf.tensor_scatter_nd_update(tensor, indices, updates))

Upvotes: 3

Views: 675

Answers (1)

Innat
Innat

Reputation: 17229

For indices, the index depth is the size or length of the index vectors. For example:

indicesA = [[1], [3], [4], [7]] # index vector with 1 element: index_depth = 1
indicesB = [[0, 1], [2, 0]]     # index vector with 2 element: index_depth = 2

The reason for indices is 2D is to hold two information, one is the length of the updates (num_updates) and the length of the index vector. Two things need to be fulfilled:

  • The index depth of indices must equal the rank of the input tensor
  • The length of updates must equal the length of the indices

So, in the example code

# tf.rank(tensor) == 1
tensor = [0, 0, 0, 0, 0, 0, 0, 0]    

# num_updates == 4, index_depth == 1 | tf.rank(indices).numpy() == 2 
indices = [[1], [3], [4], [7]]    

# num_updates == 4 | tf.rank(output).numpy() == 1  
updates = [9, 10, 11, 12]        

output = tf.tensor_scatter_nd_update(tensor, indices, updates)
tf.Tensor([ 0  9  0 10 11  0  0 12], shape=(8,), dtype=int32)

Also

# tf.rank(tensor) == 2
tensor = [[1, 1], [1, 1], [1, 1]]    

 # num_updates == 2, index_depth == 2 | tf.rank(indices).numpy() == 2
indices = [[0, 1], [2, 0]]          

# num_updates == 2 | tf.rank(output).numpy() == 2
updates = [5, 10]       
             
output = tf.tensor_scatter_nd_update(tensor, indices, updates)
tf.Tensor(
[[ 1  5]
 [ 1  1]
 [10  1]], shape=(3, 2), dtype=int32)

num_updates, index_depth = tf.convert_to_tensor(indices).shape.as_list()
[num_updates, index_depth]
[2, 2]

Upvotes: 1

Related Questions