#NEWLINE# @ opt . register_specialize ( 'fast_compile_gpu' ) #NEWLINE# @ gof . local_optimizer ( [ softmax_op ] ) #NEWLINE# def local_softmax_with_bias( node ) : #NEWLINE# #INDENT# '\n    Try to turn softmax(sum_of_stuff) -> softmax_w_bias(matrix, bias).\n\n    ' #NEWLINE# if ( node . op == softmax_op ) : #NEWLINE# #INDENT# ( x , ) = node . inputs #NEWLINE# if ( x . owner and ( x . owner . op == tensor . add ) ) : #NEWLINE# #INDENT# vectors = [ ] #NEWLINE# non_vectors = [ ] #NEWLINE# for x_in in x . owner . inputs : #NEWLINE# #INDENT# if ( list ( x_in . type . broadcastable ) == [ True , False ] ) : #NEWLINE# #INDENT# if ( x_in . owner and isinstance ( x_in . owner . op , tensor . DimShuffle ) and ( list ( x_in . owner . inputs [ 0 ] . type . broadcastable ) == [ False ] ) ) : #NEWLINE# #INDENT# vectors . append ( x_in . owner . inputs [ 0 ] ) #NEWLINE# #UNINDENT# else : #NEWLINE# #INDENT# vectors . append ( tensor . DimShuffle ( ( True , False ) , ( 1 , ) ) ( x_in ) ) #NEWLINE# #UNINDENT# #UNINDENT# else : #NEWLINE# #INDENT# non_vectors . append ( x_in ) #NEWLINE# #UNINDENT# #UNINDENT# if ( len ( non_vectors ) == 0 ) : #NEWLINE# #INDENT# assert ( len ( vectors ) > 0 ) #NEWLINE# promoted_vector = vectors . pop ( ) #NEWLINE# non_vectors . append ( tensor . shape_padleft ( promoted_vector ) ) #NEWLINE# #UNINDENT# assert non_vectors #NEWLINE# if vectors : #NEWLINE# #INDENT# if ( len ( vectors ) > 1 ) : #NEWLINE# #INDENT# vector_sum = tensor . add ( * vectors ) #NEWLINE# copy_stack_trace ( x_in , vector_sum ) #NEWLINE# #UNINDENT# else : #NEWLINE# #INDENT# vector_sum = vectors [ 0 ] #NEWLINE# #UNINDENT# if ( len ( non_vectors ) > 1 ) : #NEWLINE# #INDENT# non_vector_sum = tensor . add ( * non_vectors ) #NEWLINE# copy_stack_trace ( x_in , non_vector_sum ) #NEWLINE# #UNINDENT# else : #NEWLINE# #INDENT# non_vector_sum = non_vectors [ 0 ] #NEWLINE# #UNINDENT# try : #NEWLINE# #INDENT# sm_bias = softmax_with_bias ( non_vector_sum , node ) #NEWLINE# copy_stack_trace ( node . outputs [ 0 ] , sm_bias ) #NEWLINE# #UNINDENT# except Exception : #NEWLINE# #INDENT# return #NEWLINE# #UNINDENT# if ( sm_bias . type == node . outputs [ 0 ] . type ) : #NEWLINE# #INDENT# return [ sm_bias ]