Keras: How to implement reordering for o/p of a layer in keras?

Question

I have nine, 500 dimensional vectors as o/p from a combination of merge layer followed a reshape layer:

...
N=3
merge_rf=merge(cells_rf,mode='concat')
rf_top = Reshape((N*N, 500))(merge_rf)
#proper reordering code here - TODO

...

>>> rf_top.shape 
TensorShape([Dimension(None), Dimension(9), Dimension(500)])

I need to reorder these (9, 500) dim vectors in rf_top using a known mapping

eg - (0,1,2,3,4,5,6,7,8) to -> (0,1,2,5,4,3,6,7,8) for rf_top

Which keras layer should I use for this? and How to do it?

Answer 1

Try using:

reorder_top = merge([Lambda(lambda x: x[:, index, :], 
    output_shape = (1, 500))(rf_top) for index in permutation], 
    mode='concat', concat_axis=1)

Where permutation is the permutation according to which you want to permute your layer. The output of this layer is flattened so you should reshape it by:

reorder_top = Reshape ((N*N, 500))(reorder_top)