re-enable the grad on max in some case. The new max_and_argmax optimization do the manually done optimization that disabled the grad in some case.