Complex maths functions

ivy/data_classes/array/elementwise.py

@@ -1,6 +1,6 @@
 # global
 import abc
-from typing import Optional, Union
+from typing import Optional, Union, Literal
 
 # local
 import ivy
@@ -2552,7 +2552,12 @@ def tan(self: ivy.Array, *, out: Optional[ivy.Array] = None) -> ivy.Array:
         """
         return ivy.tan(self._data, out=out)
 
-    def tanh(self: ivy.Array, *, out: Optional[ivy.Array] = None) -> ivy.Array:
+    def tanh(
+        self: ivy.Array,
+        *,
+        out: Optional[ivy.Array] = None,
+        complex_mode: Literal["split", "magnitude", "jax"] = "jax",
+    ) -> ivy.Array:
         """
         ivy.Array instance method variant of ivy.tanh. This method simply wraps the
         function, and so the docstring for ivy.tanh also applies to this method with
@@ -2566,6 +2571,8 @@ def tanh(self: ivy.Array, *, out: Optional[ivy.Array] = None) -> ivy.Array:
         out
             optional output, 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
         -------
@@ -2581,7 +2588,7 @@ def tanh(self: ivy.Array, *, out: Optional[ivy.Array] = None) -> ivy.Array:
         >>> print(y)
         ivy.array([0., 0.762, 0.964])
         """
-        return ivy.tanh(self._data, out=out)
+        return ivy.tanh(self._data, out=out, complex_mode=complex_mode)
 
     def trunc(self: ivy.Array, *, out: Optional[ivy.Array] = None) -> ivy.Array:
         """

ivy/data_classes/container/elementwise.py

@@ -1,5 +1,5 @@
 # global
-from typing import Optional, Union, List, Dict
+from typing import Optional, Union, List, Dict, Literal
 
 # local
 import ivy
@@ -8274,6 +8274,7 @@ def _static_tanh(
         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.tanh. This method simply wraps the
@@ -8299,6 +8300,8 @@ def _static_tanh(
         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
         -------
@@ -8325,6 +8328,7 @@ def _static_tanh(
             prune_unapplied=prune_unapplied,
             map_sequences=map_sequences,
             out=out,
+            complex_mode=complex_mode,
         )
 
     def tanh(
@@ -8335,6 +8339,7 @@ def tanh(
         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.tanh. This method simply wraps the
@@ -8360,6 +8365,8 @@ def tanh(
         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
         -------
@@ -8386,6 +8393,7 @@ def tanh(
             prune_unapplied=prune_unapplied,
             map_sequences=map_sequences,
             out=out,
+            complex_mode=complex_mode,
         )
 
     @staticmethod

ivy/functional/backends/jax/elementwise.py

@@ -304,6 +304,10 @@ def log10(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray:
 
 
 def log1p(x: JaxArray, /, *, out: Optional[JaxArray] = None) -> JaxArray:
+    if x.dtype in ["complex64", "complex128"]:
+        # For certain complex numbers, jnp.log1p(z) != jnp.log(1+z)
+        # I don't know why
+        return jnp.log(1 + x)
     return jnp.log1p(x)
 
 

ivy/functional/backends/paddle/activations.py

@@ -59,7 +59,7 @@ def leaky_relu(
 
 
 @with_unsupported_device_and_dtypes(
-    {"2.5.1 and below": {"cpu": ("complex128", "complex64")}}, backend_version
+    {"2.5.1 and below": {"cpu": ("bfloat16",)}}, backend_version
 )
 def gelu(
     x: paddle.Tensor,
@@ -68,18 +68,13 @@ def gelu(
     approximate: bool = False,
     out: Optional[paddle.Tensor] = None,
 ) -> paddle.Tensor:
+    if x.dtype in [paddle.complex64, paddle.complex128]:
+        return (
+            0.5
+            * x
+            * (1 + paddle_backend.tanh(0.7978845608 * (x + 0.044715 * x * x * x)))
+        )
     if x.dtype in unsupported_dtypes:
-        if paddle.is_complex(x):
-            if approximate:
-                return (
-                    0.5
-                    * x
-                    * (
-                        1
-                        + paddle_backend.tanh(0.7978845608 * (x + 0.044715 * x * x * x))
-                    )
-                )
-            return 0.5 * x * (1 + paddle_backend.erf(x / paddle_backend.sqrt(2)))
         return F.gelu(x.cast("float32"), approximate=approximate).cast(x.dtype)
     return F.gelu(x, approximate=approximate)