chore: add complex_mode and fix supported dtypes in backends

ivy-llc · Aug 23, 2023 · 10f622f · 10f622f
1 parent fe5a626
commit 10f622f
Show file tree

Hide file tree

Showing 6 changed files with 18 additions and 4 deletions.
diff --git a/ivy/functional/backends/jax/activations.py b/ivy/functional/backends/jax/activations.py
@@ -51,6 +51,7 @@ def softplus(
     beta: Optional[Union[int, float]] = None,
     threshold: Optional[Union[int, float]] = None,
     out: Optional[JaxArray] = None,
+    complex_mode="jax",
 ) -> JaxArray:
     if beta is not None and beta != 1:
         x_beta = x * beta

diff --git a/ivy/functional/backends/mxnet/activations.py b/ivy/functional/backends/mxnet/activations.py
@@ -44,6 +44,7 @@ def softplus(
     beta: Optional[Union[(int, float)]] = None,
     threshold: Optional[Union[(int, float)]] = None,
     out: Optional[None] = None,
+    complex_mode="jax",
 ) -> None:
     if beta is not None and beta != 1:
         x_beta = x * beta

diff --git a/ivy/functional/backends/numpy/activations.py b/ivy/functional/backends/numpy/activations.py
@@ -59,6 +59,7 @@ def softplus(
     beta: Optional[Union[int, float]] = None,
     threshold: Optional[Union[int, float]] = None,
     out: Optional[np.ndarray] = None,
+    complex_mode="jax",
 ) -> np.ndarray:
     if beta is not None and beta != 1:
         x_beta = x * beta

diff --git a/ivy/functional/backends/paddle/activations.py b/ivy/functional/backends/paddle/activations.py
@@ -121,6 +121,7 @@ def softplus(
     beta: Optional[Union[int, float]] = None,
     threshold: Optional[Union[int, float]] = None,
     out: Optional[paddle.Tensor] = None,
+    complex_mode="jax",
 ) -> paddle.Tensor:
     if beta is not None and beta != 1:
         x_beta = x * beta

diff --git a/ivy/functional/backends/tensorflow/activations.py b/ivy/functional/backends/tensorflow/activations.py
@@ -47,7 +47,17 @@ def softmax(
 
 
 @with_supported_dtypes(
-    {"2.13.0 and below": ("float16", "bfloat16", "float32", "float64")}, backend_version
+    {
+        "2.13.0 and below": (
+            "float16",
+            "bfloat16",
+            "float32",
+            "float64",
+            "complex64",
+            "complex128",
+        )
+    },
+    backend_version,
 )
 def softplus(
     x: Tensor,
@@ -56,6 +66,7 @@ def softplus(
     beta: Optional[Union[int, float]] = None,
     threshold: Optional[Union[int, float]] = None,
     out: Optional[Tensor] = None,
+    complex_mode="jax",
 ) -> Tensor:
     if beta is not None and beta != 1:
         x_beta = x * beta

diff --git a/ivy/functional/backends/torch/activations.py b/ivy/functional/backends/torch/activations.py
@@ -65,15 +65,14 @@ def softmax(
     *,
     axis: Optional[int] = None,
     out: Optional[torch.Tensor] = None,
+    complex_mode="jax",
 ) -> torch.Tensor:
     if axis is None:
         axis = -1
     return torch.nn.functional.softmax(x, axis)
 
 
-@with_unsupported_dtypes(
-    {"2.0.1 and below": ("complex", "float16", "bfloat16")}, backend_version
-)
+@with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, backend_version)
 def softplus(
     x: torch.Tensor,
     /,