tfa.nn sub-module to keep the implementation of tensor operation

[WindQAQ]

System information

• TensorFlow version (you are using): nightly
• TensorFlow Addons version: source
• Is it in the tf.contrib (if so, where): no
• Are you willing to contribute it (yes/no): yes
• Are you willing to maintain it going forward? (yes/no): yes

Describe the feature and the current behavior/state.

Does anyone have some thoughts on tfa.nn sub-module? IMO, it's somewhat weird to expose tensor operations like sparsemax or gelu only in tfa.activations.

Like in core TF, the real implementations of tensor operations such as relu and selu are in tf.nn.* (most of them are done in C++, and ported to python in tensorflow.python.ops.nn_ops.py). These tensor operations in tf.nn.* are widely used in tf.keras.activations or tf.keras.layers (In fact, they use tf.keras.backend.*, which is like a keras wrapper of tf.nn.*). tf.keras.activations acts more like a hub of activation functions instead of the place to put the real implementations.

Currently, the implementation of sparsemax or gelu (WIP) are in tfa.activations, and their Layer subclass calls tfa.activations.* for call method. If we have tfa.nn sub-module, then both tfa.activations and tfa.layers use tfa.nn directly, which is much more similar with core TF style.

https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/ops/nn_ops.py
https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/keras/layers/advanced_activations.py#L313
https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/keras/activations.py

[kyleabeauchamp]

Seems reasonable to me. If we were willing to break with TF tradition, we might call this module math, as it actually seems like a core math library.

[seanpmorgan]

Certainly understand the need for this, but I'm slightly hesitant seeing as it's tough to establish a criteria for what the API of these ops should look like. To be fair tfa.activations doesn't really have a well defined API either but at least they all share a common output format. The namespace nn isn't perfect IMO because for example someone may want to use a softmax in something other than a NN graph.

I personally don't have an issue with a layer calling an activation, but there are certainly things other than activations that we may want to be included in Addons.

cc @karmel @martinwicke for thoughts.

[karmel]

I would avoid diverging and going to math. NN is, for better or for worse, a TF grab-bag, so I tend to agree that having tfa.activations makes more clear the contract that these ops adhere to. I would suggest leaving as-is. When there are things other than activations that we want to add, we can classify those accordingly. That said, I can see the argument from the other side, too, so I think you're fine either way.

[WindQAQ]

I would avoid diverging and going to math. NN is, for better or for worse, a TF grab-bag, so I tend to agree that having tfa.activations makes more clear the contract that these ops adhere to. I would suggest leaving as-is. When there are things other than activations that we want to add, we can classify those accordingly. That said, I can see the argument from the other side, too, so I think you're fine either way.

Sure. Maybe just wait for some non-activation-like tensor operations added into addons, and we could find some suitable places to classify them accordingly. Many thanks for the suggestions!