Add expit function

Author: OmarManzoor

src/array_api_extra/_lib/_funcs.py

@@ -913,3 +913,29 @@ def sinc(x: Array, /, *, xp: ModuleType | None = None) -> Array:
         xp.asarray(xp.finfo(x.dtype).eps, dtype=x.dtype, device=_compat.device(x)),
     )
     return xp.sin(y) / y
+
+
+def expit(x: Array, /, *, xp: ModuleType | None = None) -> Array:
+    """
+    Return the expit function.
+
+    The expit function, also known as the logistic sigmoid function.
+    It is the inverse of the logit function.
+
+    Parameters
+    ----------
+    x : array
+        Input array.
+    xp : array_namespace, optional
+        The standard-compatible namespace for `x`. Default: infer.
+
+    Returns
+    -------
+    array
+        An array of the same shape as x. Its entries are expit of the
+        corresponding entry of x.
+    """
+    if xp is None:
+        xp = array_namespace(x)
+
+    return 1.0 / (1.0 + xp.exp(-x))

tests/test_funcs.py

@@ -9,6 +9,7 @@
 import pytest
 from hypothesis import given
 from hypothesis import strategies as st
+from scipy import special
 
 from array_api_extra import (
     apply_where,
@@ -26,6 +27,7 @@
     sinc,
 )
 from array_api_extra._lib import Backend
+from array_api_extra._lib._funcs import expit
 from array_api_extra._lib._testing import xp_assert_close, xp_assert_equal
 from array_api_extra._lib._utils._compat import device as get_device
 from array_api_extra._lib._utils._helpers import eager_shape, ndindex
@@ -1003,3 +1005,12 @@ def test_device(self, xp: ModuleType, device: Device):
 
     def test_xp(self, xp: ModuleType):
         xp_assert_equal(sinc(xp.asarray(0.0), xp=xp), xp.asarray(1.0))
+
+
+class TestExpit:
+    def test_simple(self, xp: ModuleType):
+        x = xp.asarray([2, 3, 4, 5])
+        np_x = np.asarray([2, 3, 4, 5])
+        actual = expit(x)
+        expected = special.expit(np_x)
+        xp_assert_close(actual, expected)