Optimizing scatter_nd_* for complex tensors

[zaccharieramzi]

System information

• TensorFlow version (you are using): 2.2
• Are you willing to contribute it (Yes/No): Yes

Describe the feature and the current behavior/state.

Currently, as noted in #40605 , the scatter_nd_* functions are extremely slow.

However, this is even more extreme when used on complex tensors. This colab notebook illustrates this fact.

A simple hack consisting in treating separately real and imaginary parts makes us gain a factor 20 on computation time.

I guess at least this hack could be implemented at the python level of the op, a part on which I am definitely willing to contribute.

Will this change the current api? How?

No

Who will benefit with this feature?

Everyone using complex tensors, so I would say people working in sound processing and MRI, but this is surely not exhaustive.

Any Other info.

I don't know if this is due to eager execution or not. I don't necessarily know how to best profile these kind of issues in graph mode.

I think even if this is eager-related it still deserves a fix.

[zaccharieramzi]

For the latest nightly release (2.4.0dev-20200708), it seems that this problem is not present anymore (see the same colab).

This PR solved a problem related to tf 2.3, and maybe some other issue with scatter_nd_* was solved in tf 2.3.

Closing this since it's not relevant anymore.