From 7d547cfba8ca8709f7699405c36a690aafd787dc Mon Sep 17 00:00:00 2001 From: Gabriel de Marmiesse Date: Mon, 9 Mar 2020 02:45:51 +0100 Subject: [PATCH] Exposed the pure python gelu implementation. (#1250) --- tensorflow_addons/activations/gelu.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/tensorflow_addons/activations/gelu.py b/tensorflow_addons/activations/gelu.py index e82c58a03d..d6b630fa00 100644 --- a/tensorflow_addons/activations/gelu.py +++ b/tensorflow_addons/activations/gelu.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") @@ -42,6 +43,17 @@ def gelu(x: types.TensorLike, approximate: bool = True) -> 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 _gelu_custom_op(x, approximate) + except tf.errors.NotFoundError: + options.warn_fallback("gelu") + + return _gelu_py(x, approximate) + + +def _gelu_custom_op(x, approximate): return _activation_so.ops.addons_gelu(x, approximate)