complex dtypes support for activation functions

ivy/data_classes/array/activations.py

@@ -196,6 +196,7 @@ def softplus(
         beta: Optional[Union[int, float]] = None,
         threshold: Optional[Union[int, float]] = None,
         out: Optional[ivy.Array] = None,
+        complex_mode: Literal["split", "magnitude", "jax"] = "jax",
     ) -> ivy.Array:
         """
         ivy.Array instance method variant of ivy.softplus. This method simply wraps the
@@ -212,6 +213,8 @@ def softplus(
             the threshold parameter of the softplus function.
         out
             optional output array, for writing the result to. It must have a shape
+        complex_mode
+            optional specifier for how to handle complex data types.
 
         Returns
         -------
@@ -235,7 +238,7 @@ def softplus(
         >>> print(x)
         ivy.array([1.55, 2.13, 2.13])
         """
-        return ivy.softplus(self._data, beta=beta, threshold=threshold, out=out)
+        return ivy.softplus(self._data, beta=beta, threshold=threshold, out=out, complex_mode=complex_mode)
 
     def log_softmax(
         self: ivy.Array,

ivy/data_classes/container/activations.py

@@ -660,6 +660,7 @@ def _static_softplus(
         prune_unapplied: Union[bool, ivy.Container] = False,
         map_sequences: Union[bool, ivy.Container] = False,
         out: Optional[ivy.Container] = None,
+        complex_mode: Literal["split", "magnitude", "jax"] = "jax",
     ) -> ivy.Container:
         """
         ivy.Container static method variant of ivy.softplus. This method simply wraps
@@ -688,6 +689,8 @@ def _static_softplus(
         out
             optional output container, for writing the result to. It must have a shape
             that the inputs broadcast to.
+        complex_mode
+            optional specifier for how to handle complex data types.
 
         Returns
         -------
@@ -721,6 +724,7 @@ def _static_softplus(
             prune_unapplied=prune_unapplied,
             map_sequences=map_sequences,
             out=out,
+            complex_mode=complex_mode,
         )
 
     def softplus(
@@ -734,6 +738,7 @@ def softplus(
         prune_unapplied: Union[bool, ivy.Container] = False,
         map_sequences: Union[bool, ivy.Container] = False,
         out: Optional[ivy.Container] = None,
+        complex_mode: Literal["split", "magnitude", "jax"] = "jax",
     ) -> ivy.Container:
         """
         ivy.Container instance method variant of ivy.softplus. This method simply wraps
@@ -762,6 +767,8 @@ def softplus(
         out
             optional output container, for writing the result to. It must have a shape
             that the inputs broadcast to.
+        complex_mode
+            optional specifier for how to handle complex data types.
 
         Returns
         -------
@@ -793,6 +800,7 @@ def softplus(
             prune_unapplied=prune_unapplied,
             map_sequences=map_sequences,
             out=out,
+            complex_mode=complex_mode,
         )
 
     @staticmethod

ivy/functional/frontends/jax/nn/non_linear_activations.py

@@ -291,7 +291,7 @@ def softmax(x, axis=-1, where=None, initial=None):
 @to_ivy_arrays_and_back
 def softplus(x):
     x = _type_conversion(x)
-    return ivy.softplus(x).astype(x.dtype)
+    return ivy.softplus(x, complex_mode="jax").astype(x.dtype)
 
 
 @to_ivy_arrays_and_back

ivy/functional/ivy/activations.py

@@ -498,6 +498,32 @@ def softmax(
     return current_backend(x).softmax(x, axis=axis, out=out)
 
 
+def _wrap_between(y, a):
+    """wrap y between [-a, a]"""
+    a = ivy.array(a, dtype=y.dtype)
+    a2 = ivy.array(2 * a, dtype=y.dtype)
+    zero = ivy.array(0, dtype=y.dtype)
+    rem = ivy.remainder(ivy.add(y, a), a2)
+    rem = ivy.where(rem < zero, rem + a2, rem) - a
+    return rem
+
+
+def _softplus_jax_like(
+    x: Union[ivy.Array, ivy.NativeArray],
+    /,
+    *,
+    fn_original=None,
+    beta: Optional[Union[int, float]] = None,
+    threshold: Optional[Union[int, float]] = None,
+    out: Optional[ivy.Array] = None,
+):
+    amax = ivy.relu(x)
+    print(amax)
+    x = ivy.subtract(x, ivy.multiply(amax, ivy.array(2, dtype=x.dtype)))
+    x = ivy.add(amax, ivy.log(ivy.add(1, ivy.exp(x))))
+    x = ivy.real(x) + _wrap_between(ivy.imag(x), ivy.pi).astype(x.dtype) * ivy.astype(1j, x.dtype)
+    return x
+
 @handle_exceptions
 @handle_backend_invalid
 @handle_nestable
@@ -506,17 +532,24 @@ def softmax(
 @to_native_arrays_and_back
 @handle_array_function
 @handle_device_shifting
+@handle_complex_input
 def softplus(
     x: Union[ivy.Array, ivy.NativeArray],
     /,
     *,
     beta: Optional[Union[int, float]] = None,
     threshold: Optional[Union[int, float]] = None,
     out: Optional[ivy.Array] = None,
+    complex_mode: Literal["split", "magnitude", "jax"] = "jax",
 ) -> ivy.Array:
     """
     Apply the softplus function element-wise.
 
+    If the input is complex, then by default we apply the `log(1+ exp(x))` to  each element
+    and if the `threshold` is set we check whether the output is less than the threshold
+    we set it  to the original value of the input if not we leave as it is.
+    This behaviour can be changed by specifying a different `complex_mode`.
+
     Parameters
     ----------
     x
@@ -528,6 +561,9 @@ def softplus(
     out
         optional output array, for writing the result to. It must have a shape that the
         inputs broadcast to.
+    complex_mode
+        optional specifier for how to handle complex data types. See
+        `ivy.func_wrapper.handle_complex_input` for more detail.
 
     Returns
     -------
@@ -557,6 +593,9 @@ def softplus(
     return current_backend(x).softplus(x, beta=beta, threshold=threshold, out=out)
 
 
+softplus.jax_like = _softplus_jax_like
+
+
 @handle_exceptions
 @handle_backend_invalid
 @handle_nestable

ivy/stateful/activations.py

@@ -185,11 +185,22 @@ def _forward(self, x, *, axis=None):
 
 
 class Softplus(Module):
-    def __init__(self):
+    def __init__(
+        self,
+        complex_mode: Literal["split", "magnitude", "jax"] = "jax",
+    ):
         """Apply the SOFTPLUS activation function."""
         Module.__init__(self)
+        self._complex_mode = complex_mode
 
-    def _forward(self, x, *, beta=None, threshold=None):
+    def _forward(
+        self,
+        x,
+        *,
+        beta=None,
+        threshold=None,
+        complex_mode=None,
+    ):
         """
 
         Parameters
@@ -201,14 +212,21 @@ def _forward(self, x, *, beta=None, threshold=None):
 
         threshold
              values above this revert to a linear function. Default: ``None``.
+        complex_mode
+             optional specifier for how to handle complex data types.
 
         Returns
         -------
         ret
             The outputs following the SOFTPLUS activation *[batch_shape, d]*
 
         """
-        return ivy.softplus(x, beta=beta, threshold=threshold)
+        return ivy.softplus(
+            x,
+            beta=beta,
+            threshold=threshold,
+            complex_mode=ivy.default(complex_mode, self._complex_mode)
+        )
 
 
 class Mish(Module):

ivy_tests/test_ivy/test_frontends/test_jax/test_nn/test_non_linear_activations.py

@@ -327,10 +327,10 @@ def test_jax_softmax(
 @handle_frontend_test(
     fn_tree="jax.nn.softplus",
     dtype_and_x=helpers.dtype_and_values(
-        available_dtypes=helpers.get_dtypes("float_and_integer"),
-        large_abs_safety_factor=2,
-        small_abs_safety_factor=2,
-        safety_factor_scale="linear",
+        available_dtypes=helpers.get_dtypes("float_and_complex"),
+        large_abs_safety_factor=4,
+        small_abs_safety_factor=4,
+        safety_factor_scale="log",
     ),
     test_with_out=st.just(False),
 )

ivy_tests/test_ivy/test_functional/test_nn/test_activations.py

@@ -146,10 +146,10 @@ def test_softmax(*, dtype_and_x, axis, test_flags, backend_fw, fn_name, on_devic
 @handle_test(
     fn_tree="functional.ivy.softplus",
     dtype_and_x=helpers.dtype_and_values(
-        available_dtypes=helpers.get_dtypes("valid"),
+        available_dtypes=helpers.get_dtypes("numeric"),
         min_num_dims=1,
-        large_abs_safety_factor=8,
-        small_abs_safety_factor=8,
+        large_abs_safety_factor=4,
+        small_abs_safety_factor=4,
         safety_factor_scale="log",
     ),
     beta=st.one_of(helpers.number(min_value=0.1, max_value=10), st.none()),