CODEWITHSUNDEEP
pythontensorflow
dlfksj
dlfksj

Reputation: 75

Is there a way to sort a tensor with respect to a sub array?

I want to sort a tensor with respect to a sub-array. For example, I have a following tensor:

A = tf.constant([[4, 2, 1, 7, 5], 
                 [10, 20, 30, 40, 50]])

And I want to sort this tensor A respect to the A[0, :].

The result I expect is:

A = tf.constant([[1, 2, 4, 5, 7], 
                 [30, 20, 10, 50, 40]])

I saw a similar question in stack overflow. (Python, sort array with respect to a sub-array)

But that question is about python array, and the answer doesn't apply to my problem.

Can anyone help me? Thank you.

Upvotes: 3

Views: 742

Answers (2)

Donshel
Donshel

Reputation: 614

You can use the argsort function on A[0, :] to compute the column ordering and the gather function to compute the new tensor. Cf. tensorflow.org.

import tensorflow as tf

A = tf.constant([[4, 2, 1, 7, 5],
     [10, 20, 30, 40, 50]])

tf.gather(A, tf.argsort(A[0, :]), axis=1)

Upvotes: 1

Chris
Chris

Reputation: 29742

Using tf.gather with tf.argsort:

import tensorflow as tf:

a = tf.constant([[4, 2, 1, 7, 5], 
                 [10, 20, 30, 40, 50]])

b = tf.gather(a, tf.argsort(a[0]), axis=1)
b

Output:

<tf.Tensor: id=152, shape=(2, 5), dtype=int32, numpy=
array([[ 1,  2,  4,  5,  7],
       [30, 20, 10, 50, 40]])>

Upvotes: 2

Related Questions