@opt.register_specialize('stabilize', 'fast_compile')
@gof.local_optimizer([tensor.Elemwise])
def local_logsoftmax(node):
 if (isinstance(node.op, tensor.Elemwise) and isinstance(node.op.scalar_op, scalar.basic.Log) and (len(node.inputs) == 1) and (node.inputs[0].owner is not None) and isinstance(node.inputs[0].owner.op, Softmax)):
  inVars = node.inputs[0].owner.inputs[0]
  new_op = LogSoftmax()
  ret = new_op(inVars)
  ret.tag.values_eq_approx = values_eq_approx_remove_inf
  copy_stack_trace([node.inputs[0], node.outputs[0]], ret)
  return [ret]

@ opt . register_specialize ( 'stabilize' , 'fast_compile' ) @ gof . local_optimizer ( [ tensor . Elemwise ] ) def local_logsoftmax ( node ) : if ( isinstance ( node . op , tensor . Elemwise ) and isinstance ( node . op . scalar_op , scalar . basic . Log ) and ( len ( node . inputs ) == 1 ) and ( node . inputs [ 0 ] . owner is not None ) and isinstance ( node . inputs [ 0 ] . owner . op , Softmax ) ) : inVars = node . inputs [ 0 ] . owner . inputs [ 0 ] new_op = LogSoftmax ( ) ret = new_op ( inVars ) ret . tag . values_eq_approx = values_eq_approx_remove_inf copy_stack_trace ( [ node . inputs [ 0 ] , node . outputs [ 0 ] ] , ret ) return [ ret ]
detect log(softmax(x)) and replace it with logsoftmax(x) note: only forward pass is affected .