gaspersic/triton
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[DOCS] Fix fused softmax example script naive softmax implementation (#178)

Browse Source
This commit is contained in:
Xiangru Lian
2021-08-02 09:37:31 -07:00
committed by GitHub
parent e8031fe61f
commit 9967e9d4b4
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
python/tutorials/02-fused-softmax.py
View File

@@ -25,7 +25,7 @@ def naive_softmax(x):
# read 2MN elements ; write MN elements # read 2MN elements ; write MN elements
z = x - x_max[:, None] z = x - x_max[:, None]
# read MN elements ; write MN elements # read MN elements ; write MN elements
numerator = torch.exp(x) numerator = torch.exp(z)
# read MN elements ; write M elements # read MN elements ; write M elements
denominator = numerator.sum(dim=1) denominator = numerator.sum(dim=1)
# read 2MN elements ; write MN elements # read 2MN elements ; write MN elements