请用 PyTorch 手写实现 softmax 函数，并基于 softmax 的输出分布编写一个采样函数。

回答思路

1. 能正确实现 softmax 并处理高维输入（最后一维）
2. 说明数值稳定性处理（减去最大值）
3. 实现基于概率分布的采样（如多项分布或累积分布+随机数）
4. 给出可运行的代码并解释关键步骤