Is there a built-in function in Tensorflow for shuffling or permutating tensors?

Question:

What’s the best way to permutate a tensor along both axis (first rows and then columns or vice versa)? Should I define a py_func and do it using numpy or use one of tensor transformation functions like tf.slice – I don’t know if that’s possible.

To achieve this using numpy, I usually do the following:

>>> arr = np.arange(9).reshape([3,3])
>>> arr
array([[0, 1, 2],
       [3, 4, 5],
       [6, 7, 8]])
>>> col_perm = np.random.permutation(arr.shape[1])
>>> col_perm
array([2, 1, 0])
>>> row_perm = np.random.permutation(arr.shape[0])
>>> row_perm
array([2, 0, 1])
>>> shuffled_arr = arr[row_perm,:][:,col_perm]
>>> shuffled_arr
array([[8, 7, 6],
       [2, 1, 0],
       [5, 4, 3]])
Asked By: NULL

||

Answers:

What about tf.random_shuffle() combined with tensor transposition (tf.transpose())?

Answered By: P. Camilleri

here’s a quick function that does exactly the same as tf.random.shuffle() but also takes an axis dimension

def tf_shuffle_axis(value, axis=0, seed=None, name=None):
    perm = list(range(tf.rank(value)))
    perm[axis], perm[0] = perm[0], perm[axis]
    value = tf.random.shuffle(tf.transpose(value, perm=perm))
    value = tf.transpose(value, perm=perm)
    return value
Answered By: FarisHijazi

Though FarisHijazi’s answer is technically correct, I believe it falls short in many use cases. This is because although it shuffles along an alternative axis, it doesn’t do so uniquely.

For example

X=[[1,2,3,4],
   [1,2,3,4],
   [1,2,3,4]]

might shuffle to

X=[[3,2,4,1],
   [3,2,4,1],
   [3,2,4,1]]

using Faris’s function, but not any order in which the elements between rows are of different order.

Here is a function that shuffles along the second axis uniquely. I’ve not generalized this to any axis, but I imagine it shouldn’t be too difficult.

def tf_shuffle_second_axis(t):
    # Uniquely random along second axis
    rnd = tf.argsort(tf.random.uniform(t.shape),axis=1)
    # Add batch dimension for gathering
    rnd = tf.concat([tf.repeat(tf.range(t.shape[0])[...,tf.newaxis,tf.newaxis],tf.shape(rnd)[1],axis=1),rnd[...,tf.newaxis]],axis=2)
    # Return shuffled tensor
    return tf.gather_nd(t,rnd,batch_dims=0)

Here is an example to illustrate the difference.

t = tf.repeat(tf.random.uniform((1,4)),3,axis=0)

unique_axis_shuffle = tf_shuffle_second_axis(t)
non_unique_axis_shuffle = tf_shuffle_axis(t,axis=1)

print('Original:n{}n'.format(t))
print('Unique axis shuffle:n{}n'.format(unique_axis_shuffle))
print('Non-unique axis shuffle:n{}n'.format(non_unique_axis_shuffle))

Output:

Original:
[[0.8135692  0.10754728 0.7690911  0.84964716]
 [0.8135692  0.10754728 0.7690911  0.84964716]
 [0.8135692  0.10754728 0.7690911  0.84964716]]

Unique axis shuffle:
[[0.8135692  0.84964716 0.10754728 0.7690911 ]
 [0.7690911  0.84964716 0.8135692  0.10754728]
 [0.7690911  0.10754728 0.8135692  0.84964716]]

Non-unique axis shuffle:
[[0.10754728 0.84964716 0.8135692  0.7690911 ]
 [0.10754728 0.84964716 0.8135692  0.7690911 ]
 [0.10754728 0.84964716 0.8135692  0.7690911 ]]
Answered By: Michael Mezher
Categories: questions Tags: , ,
Answers are sorted by their score. The answer accepted by the question owner as the best is marked with
at the top-right corner.