fixes _softmax function so that nn.softmax(...) and nn.log_softmax(...) can...

fixes _softmax function so that nn.softmax(...) and nn.log_softmax(...) can accept negative axes that aren't -1 as the documentation claims. _swap() was accepting negative indices when math_ops.range() does not accept negative indices

PiperOrigin-RevId: 181771157