complex dtypes support for activation functions #21805

mohame54 · 2023-08-13T17:22:59Z

softplus activation function support for complex dtypes

github-actions · 2023-08-13T17:23:13Z

Thanks for contributing to Ivy! 😊👏
Here are some of the important points from our Contributing Guidelines 📝:
1. Feel free to ignore the run_tests (1), run_tests (2), … jobs, and only look at the display_test_results job. 👀 It contains the following two sections:
- Combined Test Results: This shows the results of all the ivy tests that ran on the PR. ✔️
- New Failures Introduced: This lists the tests that are passing on master, but fail on the PR Fork. Please try to make sure that there are no such tests. 💪
2. The lint / Check formatting / check-formatting tests check for the formatting of your code. 📜 If it fails, please check the exact error message in the logs and fix the same. ⚠️🔧
3. Finally, the test-docstrings / run-docstring-tests check for the changes made in docstrings of the functions. This may be skipped, as well. 📚
Happy coding! 🎉👨‍💻

ivy-leaves · 2023-08-13T17:30:03Z

If you are working on an open task, please edit the PR description to link to the issue you've created.

For more information, please check ToDo List Issues Guide.

Thank you 🤗

jshepherd01 · 2023-08-14T10:55:54Z

ivy/functional/ivy/activations.py

+    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


I'm not sure if I understand this implementation here, I thought softplus was simply ivy.log(ivy.add(1., ivy.exp(x))). I also can't see a beta or threshold in here, and I think we should implement at least beta for supersetting, and I think you've left a print statement in by mistake

yeah, my bad I forgot to delete the print statement, about softplus(x) = ivy.log(ivy.add(1., ivy.exp(x)) each backend has it's own implementation of softplus also some other frameworks like doesn't support complex dtype

but torch is going to make all the activation functions support support complex dtype that's why i didn't need to implement a special case for softplus op in torch.

jshepherd01 · 2023-08-14T10:56:57Z

ivy/functional/ivy/activations.py

 ) -> 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


"less than" isn't well defined for complex numbers, so you might want to be more specific here (if you decide to use threshold in the complex version at all - it doesn't seem like you have)

okay on it.

jshepherd01 · 2023-08-14T10:58:27Z

ivy/functional/ivy/activations.py

@@ -557,6 +593,9 @@ def softplus(
    return current_backend(x).softplus(x, beta=beta, threshold=threshold, out=out)


+softplus.jax_like = _softplus_jax_like


Since softplus is well defined on the complex numbers, I'd suggest that unless you're including threshold it would be better handled with changes to the backend functions rather than with a jax_like function here

rishabgit · 2023-08-17T11:17:58Z

Hey @jshepherd01, since you've already started reviewing this PR, I'm removing myself from the assignees. Thanks!
Quick note - doesn't seem to link the issue. @mohame54 hope this helps - https://unify.ai/docs/ivy/overview/contributing/the_basics.html#todo-list-issues 😄

jshepherd01

After Ved's review of my PR (#21539) I've changed the guidance I'd written on how to do these functions, so I'm also updating this review to include those changes as suggestions. There's also a small point about the dtypes in the test function.

jshepherd01 · 2023-08-17T12:23:11Z

ivy/data_classes/array/activations.py

+        complex_mode
+            optional specifier for how to handle complex data types.


Suggested change

complex_mode

optional specifier for how to handle complex data types.

complex_mode

optional specifier for how to handle complex data types. See

``ivy.func_wrapper.handle_complex_input`` for more detail.

jshepherd01 · 2023-08-17T12:23:29Z

ivy/data_classes/container/activations.py

+        complex_mode
+            optional specifier for how to handle complex data types.


Suggested change

complex_mode

optional specifier for how to handle complex data types.

complex_mode

optional specifier for how to handle complex data types. See

``ivy.func_wrapper.handle_complex_input`` for more detail.

jshepherd01 · 2023-08-17T12:23:47Z

ivy/data_classes/container/activations.py

+        complex_mode
+            optional specifier for how to handle complex data types.


Suggested change

complex_mode

optional specifier for how to handle complex data types.

complex_mode

optional specifier for how to handle complex data types. See

``ivy.func_wrapper.handle_complex_input`` for more detail.

ivy/functional/ivy/activations.py

jshepherd01 · 2023-08-17T12:25:43Z

ivy/stateful/activations.py

+    def __init__(
+        self,
+        complex_mode: Literal["split", "magnitude", "jax"] = "jax",
+    ):
        """Apply the SOFTPLUS activation function."""


Suggested change

"""Apply the SOFTPLUS activation function."""

"""

Apply the SOFTPLUS activation function.

Parameters

----------

complex_mode

Specifies how to handle complex input. See

``ivy.func_wrapper.handle_complex_input`` for more detail.

"""

Parameters should be included in docstrings

Actually, beta and threshold should really be in the __init__ part rather than _forward as well but that's beside the point of this PR

jshepherd01 · 2023-08-17T12:27:35Z

ivy/stateful/activations.py

+        complex_mode
+             optional specifier for how to handle complex data types.


Suggested change

complex_mode

optional specifier for how to handle complex data types.

We no longer use optional arguments in the forward method for stateful components, only in __init__

jshepherd01 · 2023-08-17T12:27:59Z

ivy/stateful/activations.py

+        *,
+        beta=None,
+        threshold=None,
+        complex_mode=None,


Suggested change

complex_mode=None,

jshepherd01 · 2023-08-17T12:28:25Z

ivy/stateful/activations.py

+            x,
+            beta=beta,
+            threshold=threshold,
+            complex_mode=ivy.default(complex_mode, self._complex_mode)


Suggested change

complex_mode=ivy.default(complex_mode, self._complex_mode)

complex_mode=self._complex_mode,

jshepherd01 · 2023-08-17T12:29:08Z

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

-        large_abs_safety_factor=2,
-        small_abs_safety_factor=2,
-        safety_factor_scale="linear",
+        available_dtypes=helpers.get_dtypes("float_and_complex"),


Suggested change

available_dtypes=helpers.get_dtypes("float_and_complex"),

available_dtypes=helpers.get_dtypes("numeric"),

"float_and_complex" doesn't include integers, so we should use "numeric" (or "valid") instead

jshepherd01 · 2023-08-17T12:30:20Z

ivy/functional/ivy/activations.py

+    If the input is complex, then by default we apply the `log(1+ exp(x))` to  each element0
+    This behaviour can be changed by specifying a different `complex_mode`.
+


Suggested change

If the input is complex, then by default we apply the `log(1+ exp(x))` to each element0

This behaviour can be changed by specifying a different `complex_mode`.

This behaviour is the same as with real numbers, so you could probably drop this part from the docstring - I don't think it needs to be explicitly spelled out, and we already have a description of what the complex_mode parameter does further down

jshepherd01 · 2023-08-18T13:20:38Z

ivy/functional/ivy/activations.py

+        res = ivy.where(
+            (
+                ivy.logical_or(
+                    ivy.real(x_beta) < 0, ivy.logical_and(ivy.real(x_beta) == 0, ivy.imag(x_beta) < 0)
+                )
+            ),
+            res,
+            x,
+        ).astype(x.dtype)


The actual value of threshold doesn't seem to be used here

…relu, gelu

jshepherd01 · 2023-08-21T09:23:46Z

ivy/functional/ivy/activations.py

+    if threshold is not None:
+        res = ivy.where(
+            (
+                ivy.logical_or(
+                    ivy.real(x_beta) < threshold,
+                    ivy.logical_and(
+                        ivy.real(x_beta) == threshold, ivy.imag(x_beta) < threshold
+                    ),
+                )
+            ),
+            res,
+            x,
+        ).astype(x.dtype)


Suggested change

if threshold is not None:

res = ivy.where(

(

ivy.logical_or(

ivy.real(x_beta) < threshold,

ivy.logical_and(

ivy.real(x_beta) == threshold, ivy.imag(x_beta) < threshold

),

)

),

res,

x,

).astype(x.dtype)

if threshold is not None:

res = ivy.where(

ivy.real(x_beta) < threshold,

res,

x,

).astype(x.dtype)

I've investigated a bit and I think this is a better solution, given what threshold is trying to do numerically

jshepherd01 · 2023-08-21T09:26:23Z

ivy/functional/ivy/activations.py

+    If the input is complex, then by default we apply the softplus operation
+    `log(1+ exp(x))` to  each element
+    If threshold is set we check if either its real part is strictly negative or
+    if its real part is zero and its imaginary part is negative then we apply
+    `input×β > threshold`.
+


I find this docstring quite confusing, could you reword it a bit? It still seems to reference the number being negative.

Although, I suppose instead you could remove this and edit the description of the threshold parameter (since this just describes the standard way softplus works, and then describes special behaviour of threshold with complex inputs)

jshepherd01 · 2023-08-21T11:00:15Z

ivy/functional/ivy/activations.py

+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,
+):
+    if beta is not None:
+        x_beta = ivy.multiply(x, ivy.array(beta, dtype=x.dtype))
+    else:
+        x_beta = x
+    amax = ivy.relu(x_beta)
+    res = ivy.subtract(x_beta, ivy.multiply(amax, ivy.array(2, dtype=x.dtype)))
+    res = ivy.add(amax, ivy.log(ivy.add(1, ivy.exp(res))))
+    res = ivy.real(res) + _wrap_between(ivy.imag(res), ivy.pi).astype(
+        x.dtype
+    ) * ivy.astype(1j, x.dtype)
+    if beta is not None:
+        res = ivy.divide(res, ivy.array(beta, dtype=x.dtype))
+    if threshold is not None:
+        res = ivy.where(
+            ivy.real(x_beta) < threshold,
+            res,
+            x,
+        ).astype(x.dtype)
+    return res


I'm still not sure whether this should be here or in the backends so I'm going to request a review from @Ishticode for a second opinion

Added softplus activation function support for complex dtypes

89d6de3

mohame54 requested a review from jshepherd01 August 13, 2023 17:23

mohame54 requested a review from Ishticode August 13, 2023 17:23

ivy-leaves added Array API Conform to the Array API Standard, created by The Consortium for Python Data API Standards JAX Frontend Developing the JAX Frontend, checklist triggered by commenting add_frontend_checklist Ivy Functional API labels Aug 13, 2023

ivy-seed assigned rishabgit Aug 13, 2023

jshepherd01 suggested changes Aug 14, 2023

View reviewed changes

mohame54 requested review from jshepherd01 and removed request for Ishticode August 14, 2023 18:00

Mohammed and others added 10 commits August 15, 2023 01:08

errors fixed

7501ae5

Merge branch 'unifyai:main' into main

47aa912

Merge branch 'main' of https://github.com/mohame54/ivy

28723e4

Merge branch 'unifyai:main' into main

62e95a6

Merge branch 'unifyai:main' into main

b0a6e4f

Merge branch 'unifyai:main' into main

facfc54

Merge branch 'unifyai:main' into main

a7ab849

Merge branch 'unifyai:main' into main

258d4c5

Merge branch 'unifyai:main' into main

87c63a2

Merge branch 'unifyai:main' into main

5a8dc2d

rishabgit assigned jshepherd01 and unassigned rishabgit Aug 17, 2023

jshepherd01 suggested changes Aug 17, 2023

View reviewed changes

Made softplus activation function support complex dtype

0bf323f

jshepherd01 suggested changes Aug 18, 2023

View reviewed changes

PR refactored

10c1f1a

mohame54 requested a review from jshepherd01 August 18, 2023 14:59

Mohammed and others added 8 commits August 18, 2023 19:01

refactored the some activation functions doc string like relu, leaky_…

be21aba

…relu, gelu

conflicts resolved

b742d30

lint error fixed

6060619

Merge branch 'unifyai:main' into main

25d5606

Merge branch 'main' of https://github.com/mohame54/ivy

b6a72ab

Merge branch 'unifyai:main' into main

dfc9613

Merge branch 'unifyai:main' into main

bbad132

Quick fix to the docstring of ivy.softplus

b81f2da

jshepherd01 suggested changes Aug 21, 2023

View reviewed changes

refactored threshold parameter and it's doc string

eaeaa13

mohame54 requested a review from jshepherd01 August 21, 2023 10:41

jshepherd01 reviewed Aug 21, 2023

View reviewed changes

jshepherd01 requested a review from Ishticode August 21, 2023 11:00

jshepherd01 approved these changes Aug 21, 2023

View reviewed changes

mohame54 requested review from jshepherd01 and Ishticode and removed request for Ishticode August 21, 2023 12:05

mohame54 and others added 2 commits August 22, 2023 15:44

Merge branch 'unifyai:main' into main

fe5a626

chore: add complex_mode and fix supported dtypes in backends

a5283e3

jshepherd01 force-pushed the main branch from 10f622f to a5283e3 Compare August 23, 2023 13:08

jshepherd01 merged commit 023cf84 into ivy-llc:main Aug 23, 2023
98 of 132 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

complex dtypes support for activation functions #21805

complex dtypes support for activation functions #21805

mohame54 commented Aug 13, 2023

github-actions bot commented Aug 13, 2023

ivy-leaves commented Aug 13, 2023

jshepherd01 Aug 14, 2023

mohame54 Aug 14, 2023 •

edited

Loading

mohame54 Aug 14, 2023 •

edited

Loading

jshepherd01 Aug 14, 2023

mohame54 Aug 14, 2023

jshepherd01 Aug 14, 2023

rishabgit commented Aug 17, 2023 •

edited

Loading

jshepherd01 left a comment

jshepherd01 Aug 17, 2023

jshepherd01 Aug 17, 2023

jshepherd01 Aug 17, 2023

jshepherd01 Aug 17, 2023

jshepherd01 Aug 17, 2023

jshepherd01 Aug 17, 2023

jshepherd01 Aug 17, 2023

jshepherd01 Aug 17, 2023

jshepherd01 Aug 17, 2023

jshepherd01 Aug 17, 2023

jshepherd01 Aug 18, 2023

jshepherd01 Aug 21, 2023

jshepherd01 Aug 21, 2023

jshepherd01 Aug 21, 2023

		@@ -557,6 +593,9 @@ def softplus(
		return current_backend(x).softplus(x, beta=beta, threshold=threshold, out=out)


		softplus.jax_like = _softplus_jax_like

		complex_mode
		optional specifier for how to handle complex data types.

	complex_mode=ivy.default(complex_mode, self._complex_mode)
	complex_mode=self._complex_mode,

	available_dtypes=helpers.get_dtypes("float_and_complex"),
	available_dtypes=helpers.get_dtypes("numeric"),

		If the input is complex, then by default we apply the `log(1+ exp(x))` to each element0
		This behaviour can be changed by specifying a different `complex_mode`.

complex dtypes support for activation functions #21805

complex dtypes support for activation functions #21805

Conversation

mohame54 commented Aug 13, 2023

github-actions bot commented Aug 13, 2023

ivy-leaves commented Aug 13, 2023

Choose a reason for hiding this comment

mohame54 Aug 14, 2023 • edited Loading

Choose a reason for hiding this comment

mohame54 Aug 14, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

rishabgit commented Aug 17, 2023 • edited Loading

jshepherd01 left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

mohame54 Aug 14, 2023 •

edited

Loading

mohame54 Aug 14, 2023 •

edited

Loading

rishabgit commented Aug 17, 2023 •

edited

Loading