Search code examples
pythontheano

Flip second dimension of tensor


I want to flip order of elements in the second dimension of a tensor:

x = T.tensor3('x')
f = theano.function([x], ?)
print(f(x_data))

input:

x_data = [[[1, 0, 0, 0], [0, 2, 0, 0], [0, 0, 3, 0], [0, 0, 0, 4]],
          [[5, 0, 0, 0], [0, 6, 0, 0], [0, 0, 7, 0], [0, 0, 0, 8]],
          [[9, 0, 0, 0], [0, 10, 0, 0], [0, 0, 11, 0], [0, 0, 0, 12]]
         ]

desired output:

x_data = [[[0, 0, 0, 4], [0, 0, 3, 0], [0, 2, 0, 0], [1, 0, 0, 0]],
          [[0, 0, 0, 8], [0, 0, 7, 0], [0, 6, 0, 0], [5, 0, 0, 0]],
          [[0, 0, 0, 12], [0, 0, 11, 0], [0, 10, 0, 0], [9, 0, 0, 0]]
         ]

x_data[::-1] flips the overall second dimension (not desirable):

x_data = [[[ 11.   0.   0.   0.]
           [  0.  12.   0.   0.]
           [  0.   0.  13.   0.]
           [  0.   0.   0.  14.]]

           [[  5.   0.   0.   0.]
            [  0.   6.   0.   0.]
            [  0.   0.   7.   0.]
            [  0.   0.   0.   8.]]

           [[  1.   0.   0.   0.]
            [  0.   2.   0.   0.]
            [  0.   0.   3.   0.]
            [  0.   0.   0.   4.]]]

What is the simplest way to achieve the desired output ?


Solution

  • You simply flip the dimensions you want and use full slice on the dimension before that you don't want to be changed: x_data[::, ::-1]

    import numpy as np
    x = T.tensor3('x')
    x_data = np.asarray([[[1, 0, 0, 0], [0, 2, 0, 0], [0, 0, 3, 0], [0, 0, 0, 4]],
                     [[5, 0, 0, 0], [0, 6, 0, 0], [0, 0, 7, 0], [0, 0, 0, 8]],
                     [[9, 0, 0, 0], [0, 10, 0, 0], [0, 0, 11, 0], [0, 0, 0, 12]]
                     ], dtype=theano.config.floatX)
    f = theano.function([x], x[::, ::-1])
    print(f(x_data))