[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:
@@ -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)
|
Reference in New Issue
Block a user