Optimizing models that use TensorFlow Addons activations, layers, etc

[willbattel]

System information

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

Motivation

TensorFlow Addons adds many valuable features not present in the core TensorFlow library. I have not been able to find documentation suggesting that TFAddons does or does not work with the Model Optimization library. My team wants to be able to quantize and prune models that contain TFAddons features such as the Mish activation function.

Describe the feature
This library should support the TFAddons library, if it does not already do so. Because the TFAddons functions are standard Python kernels, I think this should already be possible- so perhaps all that is needed is some documentation demonstrating this capability.

Describe how the feature helps achieve the use case
TFAddons enables expanded capability over the standard TF library. Being able to optimize models exploiting this expanded capability is important for engineers designing production-ready models.

Describe how existing APIs don't satisfy your use case (optional if obvious)
The existing APIs may already suffice, and if this is the case then all that is needed is an expansion to the documentation. If not, then a future release of this library should incorporate such capability.

[liufengdb]

@willbattel If I understand correclty, tf-addon provides a way to add custom ops, etc via some Python api?

[willbattel]

@liufengdb Take a look at this PR where they introduced the Mish activation, for example. It looks like they are adding pure Python kernels that installs into the TFA library.

The only way I can think of including Mish in an optimized model with tfmot would be with a Lambda layer. I'm wondering if there is a better way to integrate the layers from TFA.

[willbattel]

Can we get some official documentation about using TFA features with TFMOT, since they're both official TensorFlow libraries that developers will often want to use together?