Implement numerically stable batched softmax in PyTorch by shifting logits before exponentiation and normalization.
