T1-1-19 Add alpha_dropout operator

Author: GordonYang1

src/ntops/kernels/__init__.py

@@ -2,6 +2,7 @@
     abs,
     add,
     addmm,
+    alpha_dropout,
     avg_pool2d,
     bitwise_and,
     bitwise_not,
@@ -47,6 +48,7 @@
     "abs",
     "add",
     "addmm",
+    "alpha_dropout",
     "avg_pool2d",
     "bitwise_and",
     "bitwise_not",

src/ntops/kernels/alpha_dropout.py

@@ -0,0 +1,28 @@
+import functools
+
+import ninetoothed
+import ninetoothed.language as ntl
+from ninetoothed import Tensor
+
+from ntops.kernels.element_wise import arrangement
+
+
+def application(input, a, b, sat, p, seed, output):
+    keep = ntl.rand(seed, input.offsets()) > p
+    output = ntl.where(keep, a * input + b, sat)  # noqa: F841
+
+
+def premake(ndim, dtype=None, block_size=None):
+    arrangement_ = functools.partial(arrangement, block_size=block_size)
+
+    tensors = (
+        Tensor(ndim, dtype=dtype),
+        Tensor(0, dtype=ninetoothed.float64),
+        Tensor(0, dtype=ninetoothed.float64),
+        Tensor(0, dtype=ninetoothed.float64),
+        Tensor(0, dtype=ninetoothed.float64),
+        Tensor(0, dtype=ninetoothed.int64),
+        Tensor(ndim, dtype=dtype),
+    )
+
+    return arrangement_, application, tensors

src/ntops/torch/__init__.py

@@ -1,6 +1,7 @@
 from ntops.torch.abs import abs
 from ntops.torch.add import add
 from ntops.torch.addmm import addmm
+from ntops.torch.alpha_dropout import alpha_dropout
 from ntops.torch.avg_pool2d import avg_pool2d
 from ntops.torch.bitwise_and import bitwise_and
 from ntops.torch.bitwise_not import bitwise_not
@@ -46,6 +47,7 @@
     "abs",
     "add",
     "addmm",
+    "alpha_dropout",
     "avg_pool2d",
     "bitwise_and",
     "bitwise_not",

src/ntops/torch/alpha_dropout.py

@@ -0,0 +1,36 @@
+import math
+import random
+
+import torch
+
+import ntops
+from ntops.torch.utils import _cached_make
+
+# SELU saturation value: -lambda * alpha
+_ALPHA_P = -1.7580993408473766
+
+
+def alpha_dropout(input, p=0.5, training=False, inplace=False):
+    if not training or p == 0:
+        if inplace:
+            return input
+        else:
+            return input.clone()
+
+    q = 1.0 - p
+    a = 1.0 / math.sqrt(q * (1.0 + p * _ALPHA_P * _ALPHA_P))
+    b = -a * p * _ALPHA_P
+    sat = a * _ALPHA_P + b
+
+    seed = random.randrange(0, 2**31)
+
+    if inplace:
+        output = input
+    else:
+        output = torch.empty_like(input)
+
+    kernel = _cached_make(ntops.kernels.alpha_dropout.premake, input.ndim)
+
+    kernel(input, a, b, sat, p, seed, output)
+
+    return output

tests/test_alpha_dropout.py

@@ -0,0 +1,64 @@
+import math
+import random
+
+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
+
+_ALPHA_P = -1.7580993408473766
+
+
+@skip_if_cuda_not_available
+@pytest.mark.parametrize(*generate_arguments())
+def test_alpha_dropout(shape, dtype, device, rtol, atol):
+    input = torch.randn(shape, dtype=dtype, device=device)
+    p = random.uniform(0.1, 0.5)
+
+    ninetoothed_output = ntops.torch.alpha_dropout(input, p=p, training=True)
+    reference_output = F.alpha_dropout(input, p=p, training=True)
+
+    # 1. Shape must match.
+    assert ninetoothed_output.shape == reference_output.shape
+
+    # 2. Compute expected affine parameters.
+    q = 1.0 - p
+    a = 1.0 / math.sqrt(q * (1.0 + p * _ALPHA_P * _ALPHA_P))
+    b = -a * p * _ALPHA_P
+    sat = a * _ALPHA_P + b
+
+    # 3. Drop ratios should be close to each other.
+    ninetoothed_drop_ratio = (
+        torch.isclose(
+            ninetoothed_output, torch.full_like(ninetoothed_output, sat), atol=atol
+        )
+        .float()
+        .mean()
+        .item()
+    )
+    reference_drop_ratio = (
+        torch.isclose(
+            reference_output, torch.full_like(reference_output, sat), atol=atol
+        )
+        .float()
+        .mean()
+        .item()
+    )
+
+    assert abs(ninetoothed_drop_ratio - reference_drop_ratio) < 0.1
+
+    # 4. Kept elements should satisfy the same affine transform.
+    kept_mask = ~torch.isclose(
+        ninetoothed_output, torch.full_like(ninetoothed_output, sat), atol=atol
+    )
+    expected_kept = a * input[kept_mask].float() + b
+    actual_kept = ninetoothed_output[kept_mask].float()
+
+    assert torch.allclose(actual_kept, expected_kept, rtol=rtol, atol=atol)
+
+    # 5. training=False should return input unchanged.
+    output_eval = ntops.torch.alpha_dropout(input, p=p, training=False)
+    assert torch.equal(output_eval, input)