Add gelu operator

Author: voltjia

src/ntops/__init__.py

@@ -3,6 +3,7 @@
 from ntops.addmm import addmm
 from ntops.bmm import bmm
 from ntops.div import div
+from ntops.gelu import gelu
 from ntops.mm import mm
 
-__all__ = ["abs", "add", "addmm", "bmm", "div", "mm"]
+__all__ = ["abs", "add", "addmm", "bmm", "div", "gelu", "mm"]

src/ntops/gelu.py

@@ -0,0 +1,50 @@
+import functools
+import math
+
+import ninetoothed
+import ninetoothed.language as ntl
+import torch
+from ninetoothed import Tensor
+
+from ntops import element_wise
+
+
+def default_application(input, output):
+    output = input * 0.5 * (1 + ntl.erf(input / ntl.sqrt(2.0)))  # noqa: F841
+
+
+def tanh_application(input, output):
+    input_loaded = input
+
+    output = (  # noqa: F841
+        0.5
+        * input_loaded
+        * (
+            1
+            + ntl.tanh(
+                ntl.sqrt(2 / math.pi) * (input_loaded + 0.044715 * input_loaded**3)
+            )
+        )
+    )
+
+
+def gelu(input, approximate="none"):
+    output = torch.empty_like(input)
+
+    kernel = _make(input.ndim, approximate)
+
+    kernel(input, output)
+
+    return output
+
+
+@functools.cache
+def _make(ndim, approximate):
+    tensors = (Tensor(ndim), Tensor(ndim))
+
+    if approximate == "tanh":
+        application = tanh_application
+    else:
+        application = default_application
+
+    return ninetoothed.make(element_wise.arrangement, application, tensors)

tests/test_gelu.py

@@ -0,0 +1,23 @@
+import pytest
+import torch
+import torch.nn.functional as F
+
+import ntops
+from tests.skippers import skip_if_cuda_not_available
+from tests.utils import generate_arguments
+
+
+@skip_if_cuda_not_available
+@pytest.mark.parametrize(*generate_arguments())
+def test_cuda(shape, dtype, atol, rtol):
+    device = "cuda"
+
+    input = torch.randn(shape, dtype=dtype, device=device)
+
+    for approximate in ("none", "tanh"):
+        ninetoothed_output = ntops.gelu(input)
+        reference_output = F.gelu(input)
+
+        assert torch.allclose(
+            ninetoothed_output, reference_output, atol=atol, rtol=rtol
+        )