Rewritten argmax, softmax kernel; bug fixes