Exposed the pure python implementation of mish and softshrink. (#1252)

* Exposed the pure python implementation of mish.

* Exposed the pure implementation of softshrink publicly.

* Fix name.

* The same error is not raise in python and in the custom op.

tensorflow_addons/activations/mish.py

@@ -17,6 +17,7 @@
 
 from tensorflow_addons.utils import types
 from tensorflow_addons.utils.resource_loader import LazySO
+from tensorflow_addons import options
 
 _activation_so = LazySO("custom_ops/activations/_activation_ops.so")
 
@@ -36,6 +37,17 @@ def mish(x: types.TensorLike) -> tf.Tensor:
         A `Tensor`. Has the same type as `x`.
     """
     x = tf.convert_to_tensor(x)
+
+    if not options.TF_ADDONS_PY_OPS:
+        try:
+            return _mish_custom_op(x)
+        except tf.errors.NotFoundError:
+            options.warn_fallback("mish")
+
+    return _mish_custom_op(x)
+
+
+def _mish_custom_op(x):
     return _activation_so.ops.addons_mish(x)
 
 

tensorflow_addons/activations/softshrink.py

@@ -18,6 +18,7 @@
 
 from tensorflow_addons.utils import types
 from tensorflow_addons.utils.resource_loader import LazySO
+from tensorflow_addons import options
 
 _activation_so = LazySO("custom_ops/activations/_activation_ops.so")
 
@@ -40,6 +41,17 @@ def softshrink(
         A `Tensor`. Has the same type as `x`.
     """
     x = tf.convert_to_tensor(x)
+
+    if not options.TF_ADDONS_PY_OPS:
+        try:
+            return _softshrink_custom_op(x, lower, upper)
+        except tf.errors.NotFoundError:
+            options.warn_fallback("softshrink")
+
+    return _softshrink_py(x, lower, upper)
+
+
+def _softshrink_custom_op(x, lower, upper):
     return _activation_so.ops.addons_softshrink(x, lower, upper)
 
 

tensorflow_addons/activations/softshrink_test.py

@@ -18,15 +18,18 @@
 import numpy as np
 import tensorflow as tf
 from tensorflow_addons.activations import softshrink
-from tensorflow_addons.activations.softshrink import _softshrink_py
+from tensorflow_addons.activations.softshrink import (
+    _softshrink_py,
+    _softshrink_custom_op,
+)
 from tensorflow_addons.utils import test_utils
 
 
 @test_utils.run_all_in_graph_and_eager_modes
 class SoftshrinkTest(tf.test.TestCase, parameterized.TestCase):
     def test_invalid(self):
         with self.assertRaisesOpError("lower must be less than or equal to upper."):
-            y = softshrink(tf.ones(shape=(1, 2, 3)), lower=2.0, upper=-2.0)
+            y = _softshrink_custom_op(tf.ones(shape=(1, 2, 3)), lower=2.0, upper=-2.0)
             self.evaluate(y)
 
     @parameterized.named_parameters(