@gof.optimizer
def crossentropy_to_crossentropy_with_softmax(fgraph):
 def search_make_one_sub():
  for node in fgraph.toposort():
   if (node.op == crossentropy_categorical_1hot):
    (nll,) = node.outputs
    (sm, one_of_n) = node.inputs
    if (sm.owner and (sm.owner.op == softmax_op)):
     (x,) = sm.owner.inputs
     (new_nll, new_sm, new_am) = crossentropy_softmax_argmax_1hot_with_bias(x, tensor.zeros_like(x[0]), one_of_n)
     fgraph.replace_all_validate([(nll, new_nll), (sm, new_sm)], reason='crossentropy_to_crossentropy_with_softmax')
     return True
    if (sm.owner and (sm.owner.op == softmax_with_bias)):
     (x, b) = sm.owner.inputs
     (new_nll, new_sm, new_am) = crossentropy_softmax_argmax_1hot_with_bias(x, b, one_of_n)
     fgraph.replace_all_validate([(nll, new_nll), (sm, new_sm)], reason='crossentropy_to_crossentropy_with_softmax')
     return True
  return False
 while search_make_one_sub():
  pass
 return

@ gof . optimizer def crossentropy_to_crossentropy_with_softmax ( fgraph ) : def search_make_one_sub ( ) : for node in fgraph . toposort ( ) : if ( node . op == crossentropy_categorical_1hot ) : ( nll , ) = node . outputs ( sm , one_of_n ) = node . inputs if ( sm . owner and ( sm . owner . op == softmax_op ) ) : ( x , ) = sm . owner . inputs ( new_nll , new_sm , new_am ) = crossentropy_softmax_argmax_1hot_with_bias ( x , tensor . zeros_like ( x [ 0 ] ) , one_of_n ) fgraph . replace_all_validate ( [ ( nll , new_nll ) , ( sm , new_sm ) ] , reason = 'crossentropy_to_crossentropy_with_softmax' ) return True if ( sm . owner and ( sm . owner . op == softmax_with_bias ) ) : ( x , b ) = sm . owner . inputs ( new_nll , new_sm , new_am ) = crossentropy_softmax_argmax_1hot_with_bias ( x , b , one_of_n ) fgraph . replace_all_validate ( [ ( nll , new_nll ) , ( sm , new_sm ) ] , reason = 'crossentropy_to_crossentropy_with_softmax' ) return True return False while search_make_one_sub ( ) : pass return
this is a stabilization optimization that is more general than crossentropy_to_crossentropy_with_softmax_with_bias .