#NEWLINE# def get_output_for( self , input , * args , **kwargs ) : #NEWLINE# #INDENT# bs = ( input . shape [ 0 ] // self . num_columns ) #NEWLINE# n = T . prod ( input . shape [ 1 : ] ) #NEWLINE# input = input . reshape ( ( self . num_columns , bs , n ) ) #NEWLINE# return input . dimshuffle ( 1 , 0 , 2 ) . reshape ( ( bs , ( self . num_columns * n ) ) )