add gelu_fast

HeegyuKim · Oct 6, 2023 · 194d4f0 · 194d4f0
1 parent 6238285
commit 194d4f0
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/src/transformers/modeling_flax_utils.py b/src/transformers/modeling_flax_utils.py
@@ -63,12 +63,17 @@ def quick_gelu(x):
     return x * jax.nn.sigmoid(1.702 * x)
 
 
+def gelu_fast(x):
+    return 0.5 * x * (1.0 + jnp.tanh(x * 0.7978845608 * (1.0 + 0.044715 * x * x)))
+
+
 ACT2FN = {
     "gelu": partial(nn.gelu, approximate=False),
     "relu": nn.relu,
     "silu": nn.swish,
     "swish": nn.swish,
     "gelu_new": partial(nn.gelu, approximate=True),
+    "gelu_fast": gelu_fast,
     "quick_gelu": quick_gelu,
 }