CODEWITHSUNDEEP
pythontensorflowtensorflow2.0tensor
Francis
Francis

Reputation: 199

How to reduce dimension of a tensor with keeping specific indexed elements of that dimension?

The reduce of dimension is similar to what reduce_max() does, the difference is I want a specific index of the element in that dimension instead of simply picking the maximal one. For example, I have a 2x3 tensor A = [[0,1,2],[2,2,0]]. If I apply tf.argmax(A), I get index tensor [1, 1, 0]. How can I use this index tensor [1, 1, 0] to get the tensor as tf.reduce_max(A, 0) = [2, 2, 2]?

The reason I am not using tf.reduce_max directly is I want to use different index tensor instead of the argmax index tensor to reduce the dimension or keep the indexed value instead of the maximal value at that dimension.

Upvotes: 1

Views: 1000

Answers (1)

Matus Dubrava
Matus Dubrava

Reputation: 14462

You can use tf.gather_nd function to do that but you will need to convert that [1, 1, 0] index tensor into 2D tensor.

Here I assume that the index tensor is a numpy array (you can convert tensorflow tensor into numpy array by calling .numpy() method.

idx = np.array([1, 1, 0])

idx = np.c_[idx[:, np.newaxis], np.arange(len(idx))]

print(idx)

# Output:
# array([[1, 0],
#        [1, 1],
#        [0, 2]])

Which means: pick (row1, col0), (row1, col1), and (row0, col2) when using the above mentioned tf.gather_nd

A = tf.Variable([[0, 1, 2], [2, 2, 0]])

tf.gather_nd(A, idx)

Will give you the expected [2, 2, 2] tensor.

Upvotes: 2

Related Questions