Adding a pure python implementation corresponding to simple C++/Cuda ops.

[gabrieldemarmiesse]

Describe the feature and the current behavior/state.

For custom ops, we have two implementations: C++ and Cuda. I propose that for simple ops , meaning less than 10 python lines, we implement a pure Python version too. I see several use cases:

• We can more easily test the correctness of the C++ and Cuda implementation. In the test suite, we generate random tensors and check that the results is the same for the three implementations. Since python code is easy to read, we have a readable reference implementation that's easier to check.
• It helps with users who want to understand the source code.
• It should also help all users who have issues with custom ops. It's better for those users to have a slow implementation, than an implementation which doesn't run at all. It also helps for people who want to run those ops in specialized hardware (ex: ROCm). We've had a number of issues related to users having to deal with that: #1109 #1099 #990 #987 .... We should provide a flag for them which they can activate to use the pure python implementation. For users of tensorflow addons on windows and mac, they'll be able to use their gpu.

Relevant information

• Are you willing to contribute it (yes/no): yes
• Are you willing to maintain it going forward? (yes/no): yes
• Is there a relevant academic paper? (if so, where): no
• Is there already an implementation in another framework? (if so, where): no
• Was it part of tf.contrib? (if so, where): no

Which API type would this fall under (layer, metric, optimizer, etc.)

Activation, image, layers, seq2seq, text

Who will benefit with this feature?

People having problems with custop ops, the maintainers because the test suite will be more robust.

Any other info.

We should make a pull request for a simple op see how to implement it. Especially how do we access the different implementations from the user side and tests side. Good candidates would be activations functions.

[failure-to-thrive]

Frankly, it's not clear for me why TFA activations/... need C++ code. All those functions can be expressed with regular Tensorflow ops, gaining the same benefits such as a GPU placement. What did I miss?

[gabrieldemarmiesse]

https://github.com/tensorflow/addons/pull/913 can give some background and discussion that we had concerning this.

[seanpmorgan]

I really like the idea of this and it's ability to test python vs custom-op implementations is much needed. I'm a little squeemish regarding the flag to do this though. Do you imagine it would be a flag in for example the activation API or some global setting?

[gabrieldemarmiesse]

I've no idea what kind of flag we should use for the public API. Ideas welcome :)

[georgesterpu]

@gabrieldemarmiesse what do you think of porting the Beam search decoder from tensorflow / models / official / nlp / transformer / beam_search.py to addons as a python alternative ?

[failure-to-thrive]

Also please keep mind that that you should add python implementations of gradients, as C++ do. Autodeduction of them is quite suboptimal and should be hand-crafted.

[gabrieldemarmiesse]

@georgesterpu it seems a bit too complex. I'm afraid it would actually become a maintainance burden instead of helping us. It might be better to not copy the implementation here.

[gabrieldemarmiesse]

I made pull requests to expose the implementations publicly:

1137

1250

1251

1252

1253

[bhack]

@georgesterpu I dont know if it could be better that you propose the adoption in https://github.com/tensorflow/text. So if they will get the ops in charge we could evaluate to deprecate, soon or later, the custom c++/cuda version in addons.

[guillaumekln]

For reference, #1925 is adding a Python implementation for tfa.seq2seq.gather_tree which should simplify model export.

[seanpmorgan]

TensorFlow Addons is transitioning to a minimal maintenance and release mode. New features will not be added to this repository. For more information, please see our public messaging on this decision: TensorFlow Addons Wind Down

Please consider sending feature requests / contributions to other repositories in the TF community with a similar charters to TFA: Keras Keras-CV Keras-NLP