feature_alpha_dropout

ivy/functional/frontends/torch/nn/functional/dropout_functions.py

@@ -7,23 +7,26 @@
 
 # ToDo: this function will be simplified once ivy.alpha_dropout is implemented
 @to_ivy_arrays_and_back
-@with_unsupported_dtypes({"2.0.1 and below": ("float16",)}, "torch")
+@with_unsupported_dtypes({"2.0.1 and below": ("uint16", "uint32", "uint64")}, "torch")
 def alpha_dropout(input, p=0.5, training=False, inplace=False):
     if p == 0.0 or not training or input.shape == () or input.shape == (0,):
         return input
-    neg_saturation = ivy.log1p(ivy.exp(-ivy.square(input)))
+    alpha = 1.7580993408473766
+    a = float(1.0 / ivy.sqrt((alpha * alpha * p + 1) * (1 - p)))
     mask = ivy.where(
         ivy.random_uniform(shape=input.shape, device=ivy.dev(input)) < p,
         0.0,
         1.0,
     )
+    b = ((mask - 1) * alpha * a) + alpha * a * p
+    mask *= a
+
     if inplace:
-        ivy.inplace_update(input, mask * input + (1 - mask) * neg_saturation)
-        ivy.inplace_update(input, input / ivy.sqrt(1 - p / (1 - p + 1e-5)))
+        ivy.inplace_update(input, mask * input + b)
         return input
     else:
-        masked = mask * input + (1 - mask) * neg_saturation
-        return masked / ivy.sqrt(1 - p / (1 - p + 1e-5))
+        masked = mask * input + b
+        return masked
 
 
 @to_ivy_arrays_and_back
@@ -61,3 +64,33 @@ def dropout3d(input, p=0.5, training=True, inplace=False):
             input, p, training=training, data_format="NDHWC", out=input
         )
     return ivy.dropout3d(input, p, training=training, data_format="NDHWC")
+
+
+# ToDo: this function will be simplified once ivy.feature_alpha_dropout is implemented
+@to_ivy_arrays_and_back
+@with_unsupported_dtypes({"2.0.1 and below": ("uint16", "uint32", "uint64")}, "torch")
+def feature_alpha_dropout(input, p=0.5, training=False, inplace=False):
+    if p == 0.0 or not training or len(input.shape) < 4:
+        return input
+    alpha = 1.7580993408473766
+    a = float(1.0 / ivy.sqrt((alpha * alpha * p + 1) * (1 - p)))
+
+    mask_shape = input.shape[0:2]
+    for _ in range(2, len(input.shape)):
+        mask_shape += (1,)
+    feature_mask = ivy.where(
+        ivy.random_uniform(shape=mask_shape, device=ivy.dev(input)) < p,
+        0.0,
+        1.0,
+    )
+    mask = ivy.ones(input.shape) * feature_mask
+
+    b = ((mask - 1) * alpha * a) + alpha * a * p
+    mask *= a
+
+    if inplace:
+        ivy.inplace_update(input, mask * input + b)
+        return input
+    else:
+        masked = mask * input + b
+        return masked

ivy_tests/test_ivy/test_frontends/test_torch/test_nn/test_functional/test_dropout_functions.py

@@ -238,3 +238,48 @@ def test_torch_dropout3d(
     for u in ret:
         # cardinality test
         assert u.shape == x.shape
+
+
+@handle_frontend_test(
+    fn_tree="torch.nn.functional.feature_alpha_dropout",
+    dtype_and_x=helpers.dtype_and_values(
+        available_dtypes=helpers.get_dtypes("float"),
+        min_value=0,
+        max_value=50,
+        allow_inf=False,
+        min_num_dims=4,
+        min_dim_size=2,
+    ),
+    prob=helpers.floats(min_value=0, max_value=0.9),
+    training=st.booleans(),
+    test_inplace=st.just(False),
+)
+def test_torch_feature_alpha_dropout(
+    *,
+    dtype_and_x,
+    prob,
+    training,
+    on_device,
+    fn_tree,
+    frontend,
+    test_flags,
+    backend_fw,
+):
+    input_dtype, x = dtype_and_x
+    ret = helpers.test_frontend_function(
+        input_dtypes=input_dtype,
+        backend_to_test=backend_fw,
+        frontend=frontend,
+        test_flags=test_flags,
+        fn_tree=fn_tree,
+        on_device=on_device,
+        input=x[0],
+        p=prob,
+        training=training,
+        test_values=False,
+    )
+    ret = helpers.flatten_and_to_np(ret=ret, backend=backend_fw)
+    x = np.asarray(x[0], input_dtype[0])
+    for u in ret:
+        # cardinality test
+        assert u.shape == x.shape