From ea68239ceae9e927287e0c26dcbdfc765b92de92 Mon Sep 17 00:00:00 2001 From: gabrieldemarmiesse Date: Sun, 8 Mar 2020 13:52:56 +0000 Subject: [PATCH] Exposed the pure python implementation of lisht. --- tensorflow_addons/activations/lisht.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/tensorflow_addons/activations/lisht.py b/tensorflow_addons/activations/lisht.py index a4cb7febde..7803f403e0 100644 --- a/tensorflow_addons/activations/lisht.py +++ b/tensorflow_addons/activations/lisht.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 lisht(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 _lisht_custom_op(x) + except tf.errors.NotFoundError: + options.warn_fallback("lisht") + + return _lisht_py(x) + + +def _lisht_custom_op(x): return _activation_so.ops.addons_lisht(x)