#NEWLINE# @ opt . register_specialize ( 'stabilize' , 'fast_compile' ) #NEWLINE# @ gof . local_optimizer ( [ SoftmaxGrad ] ) #NEWLINE# def local_logsoftmax_grad( node ) : #NEWLINE# #INDENT# "\n    Detect Log(Softmax(x))'s grad and replace it with LogSoftmax(x)'s grad\n\n    Note: only grad is affected\n    " #NEWLINE# if ( isinstance ( node . op , SoftmaxGrad ) and ( len ( node . inputs ) == 2 ) and ( node . inputs [ 0 ] . owner is not None ) and isinstance ( node . inputs [ 0 ] . owner . op , tensor . Elemwise ) and ( len ( node . inputs [ 0 ] . owner . inputs ) >= 2 ) and ( node . inputs [ 0 ] . owner . inputs [ 1 ] . owner is not None ) and ( node . inputs [ 0 ] . owner . inputs [ 1 ] . owner . op == softmax_op ) and ( node . inputs [ 1 ] == node . inputs [ 0 ] . owner . inputs [ 1 ] ) and ( not ( ( node . inputs [ 0 ] . owner . op == tensor . true_div ) and ( node . inputs [ 0 ] . owner . inputs [ 0 ] . owner is not None ) and isinstance ( node . inputs [ 0 ] . owner . inputs [ 0 ] . owner . op , subtensor . AdvancedIncSubtensor ) ) ) ) : #NEWLINE# #INDENT# sm = node . inputs [ 0 ] . owner . inputs [ 1 ] #NEWLINE# grads = node . inputs [ 0 ] . owner . inputs [ 0 ] #NEWLINE# if ( grads . broadcastable [ 1 ] and ( not sm . broadcastable [ 1 ] ) ) : #NEWLINE# #INDENT# grads = tensor . alloc ( grads , grads . shape [ 0 ] , sm . shape [ 1 ] ) #NEWLINE# #UNINDENT# ret = ( grads - ( tensor . sum ( grads , axis = 1 , keepdims = True ) * sm ) ) #NEWLINE# ret . tag . values_eq_approx = values_eq_approx_remove_nan #NEWLINE# return [ ret ]