forked from ccf-ai-infra/Intro-ops
24 lines
498 B
Python
24 lines
498 B
Python
from __future__ import annotations
|
|
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
sys.path.insert(0, str(ROOT / "python"))
|
|
|
|
import torch
|
|
|
|
from operator_runtime import softmax
|
|
|
|
|
|
def main() -> None:
|
|
src = torch.randn((32, 128), device="cuda", dtype=torch.float32)
|
|
out = softmax(src, dim=1, backend="nvidia")
|
|
torch.testing.assert_close(out, torch.softmax(src, dim=1), atol=1e-5, rtol=1e-5)
|
|
print("softmax ok")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|
|
|