How to rotate a tensor row-wise in tensorflow?

Question

I would like to be able to rotate a tensor in a fashion similar to collections.deque.rotate. Being able to broadcast this batch-wise would be ideal. For example, given an array of rotations foo, I would like to rotate each row of bar by its respective entry in foo.

e.g.

def tf_rotate(to_rotate, rotate_by):
    pass #How to do this?

with tf.Session('') as sesh:
    foo = tf.constant([1,2,3])
    bar = tf.constant(
        [[1,  2,  3],
         [4,  5,  6],
         [7,  8,  9]])
    print(
        sesh.run(
            tf_rotate(bar, foo)
    ))

Should give:

[[3  1  2]
 [5  6  4]
 [7  8  9]]

This operation is pretty trivial as a for loop, but such things become complicated when implementing an algorithm in tensorflow. How can we perform a circular permutation to an array/tensor in tensorflow?

Asked By: Him

||

Source

Answer 1

May be sth like this:

def rotate(matrix, shifts):
    """"requested rotate function - assumes matrix shape is mxn and shifts shape is m"""

    # get shape of the input matrix
    shape = tf.shape(matrix)

    # compute and stack the meshgrid to get the index matrix of shape (2,m,n)
    ind = tf.stack(tf.meshgrid(tf.range(shape[0]), tf.range(shape[1]), indexing='ij'))

    # reshape it to (m,n,2)
    ind = tf.transpose(ind, [1,2,0])

    # add the value from shifts to the corresponding row and devide modulo shape[1]
    # this will effectively introduce the desired shift, but at the level of indices
    shifted_ind = tf.mod(tf.transpose(tf.transpose(ind[:,:,1]) + shifts), shape[1])

    # convert the shifted indices to the right shape
    new_ind = tf.transpose(tf.stack([ind[:,:,0], shifted_ind]) , [1,2,0]) 

    # return the resliced tensor
    return tf.gather_nd(matrix, new_ind)

Example:

In [36]: matrix = np.tile(10*np.arange(10), [7,1])                                                                                        

In [37]: matrix                                                                                                                           
Out[37]: 
array([[ 0, 10, 20, 30, 40, 50, 60, 70, 80, 90],
       [ 0, 10, 20, 30, 40, 50, 60, 70, 80, 90],
       [ 0, 10, 20, 30, 40, 50, 60, 70, 80, 90],
       [ 0, 10, 20, 30, 40, 50, 60, 70, 80, 90],
       [ 0, 10, 20, 30, 40, 50, 60, 70, 80, 90],
       [ 0, 10, 20, 30, 40, 50, 60, 70, 80, 90],
       [ 0, 10, 20, 30, 40, 50, 60, 70, 80, 90]])

In [38]: shifts = np.random.randint(low=0, high=matrix.shape[1], size=matrix.shape[0], dtype=int32)                                       

In [39]: shifts                                                                                                                           
Out[39]: array([5, 7, 6, 4, 5, 6, 6], dtype=int32)

In [40]: res = rotate(tf.constant(matrix), tf.constant(shifts))                                                                           

In [41]: sess.run(res)                                                                                                                    
Out[41]: 
array([[50, 60, 70, 80, 90,  0, 10, 20, 30, 40],
       [70, 80, 90,  0, 10, 20, 30, 40, 50, 60],
       [60, 70, 80, 90,  0, 10, 20, 30, 40, 50],
       [40, 50, 60, 70, 80, 90,  0, 10, 20, 30],
       [50, 60, 70, 80, 90,  0, 10, 20, 30, 40],
       [60, 70, 80, 90,  0, 10, 20, 30, 40, 50],
       [60, 70, 80, 90,  0, 10, 20, 30, 40, 50]])

Not sure if it is better than a loop though.

Answered By: Alexandr Dibrov

Answer 2

thanks for your answer. It works great. The one issue I have is that I need to run this model on Edge TPU but it does not support the gather_nd operation. Is there a way to do this with the operators that are supported on the Edge TPU? Thanks!

Answered By: gsirocco

How to rotate a tensor row-wise in tensorflow?

Question:

Answers: