(x: HLOTensor, axis: int = -1)
| 711 | |
| 712 | |
| 713 | def softmax(x: HLOTensor, axis: int = -1): |
| 714 | assert isinstance(axis, int), f"axis should be int, but get {axis}({type(axis)})" |
| 715 | x_max = x.max(axis=axis, keepdims=True) |
| 716 | x_exp = exp(x - x_max) |
| 717 | x_exp_sum = x_exp.sum(axis=axis, keepdims=True) |
| 718 | y = x_exp / x_exp_sum |
| 719 | return y |
| 720 | |
| 721 | |
| 722 | def softmax_grad(y: HLOTensor, dy: HLOTensor, axis: int = -1): |
no test coverage detected