#NEWLINE# def get_output_for( self , input , * args , **kwargs ) : #NEWLINE# #INDENT# s = input . shape #NEWLINE# input_unfolded = input . reshape ( ( 8 , ( s [ 0 ] // 8 ) , s [ 1 ] , p [ 2 ] , s [ 3 ] ) ) #NEWLINE# permuted_inputs = [ ] #NEWLINE# for ( p , inv_tf ) in zip ( self . perm_matrix , self . inv_tf_funcs ) : #NEWLINE# #INDENT# input_permuted = inv_tf ( input_unfolded [ p ] . reshape ( s ) ) #NEWLINE# permuted_inputs . append ( input_permuted ) #NEWLINE# #UNINDENT# return nn . utils . concatenate ( permuted_inputs , axis = 1 )