[PYTHON] Removed support for dense softmax

Interest seems limited now that it is fused in cross_entropy. Will
likely re-add once it's easier to share code between ops
This commit is contained in:
Philippe Tillet
2021-02-07 16:46:47 -05:00
parent 682ac4c60e
commit b0647cfd52
3 changed files with 0 additions and 41 deletions

View File

@@ -1,8 +0,0 @@
import torch
import triton
def test_op(M = 1024, N = 1024, dtype = torch.float32):
x = torch.randn(M, N, dtype=dtype, device='cuda')
th_y = torch.softmax(x, dim=-1)
tt_y = triton.ops.softmax(x)
assert torch.allclose(tt_y, th_y)