2020 the huggingface team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license gaussian error linear unit original implementation of the gelu activation function in google bert repo when initially created for information openai gpt s gelu is slightly different and gives slightly different results 0 5 x 1 torch tanhmath sqrt2 math pi x 0 044715 torch powx 3 also see https arxiv orgabs1606 08415 gaussian error linear unit this is a smoother version of the gelu original paper https arxiv orgabs1606 0841 args x float tensor to perform activation returns x with the gelu activation applied clip the range of possible gelu outputs between 10 10 this is especially useful for quantization purpose as it allows mapping 2 negatives values in the gelu spectrum for more information on this trick please refer to https arxiv orgabs2004 09602 gaussian error linear unit original implementation of the gelu activation function in google bert repo when initially created for information openai gpt s gelu is slightly different and gives slightly different results 0 5 x 1 torch tanhmath sqrt2 math pi x 0 044715 torch powx 3 also see https arxiv orgabs1606 08415 param x return gated linear unit implementation as defined in the original paper see https arxiv orgabs1612 08083 where the input x is split in two halves across a dimension axis a and b returning a sigmoidb args x float tensor to perform activation axis dimension across which x be split in half returns x with the glu activation applied with its size halved across the dimension axis 2020 the huggingface team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license gaussian error linear unit original implementation of the gelu activation function in google bert repo when initially created for information openai gpt s gelu is slightly different and gives slightly different results 0 5 x 1 torch tanh math sqrt 2 math pi x 0 044715 torch pow x 3 also see https arxiv org abs 1606 08415 gaussian error linear unit this is a smoother version of the gelu original paper https arxiv org abs 1606 0841 args x float tensor to perform activation returns x with the gelu activation applied clip the range of possible gelu outputs between 10 10 this is especially useful for quantization purpose as it allows mapping 2 negatives values in the gelu spectrum for more information on this trick please refer to https arxiv org abs 2004 09602 gaussian error linear unit original implementation of the gelu activation function in google bert repo when initially created for information openai gpt s gelu is slightly different and gives slightly different results 0 5 x 1 torch tanh math sqrt 2 math pi x 0 044715 torch pow x 3 also see https arxiv org abs 1606 08415 param x return gated linear unit implementation as defined in the original paper see https arxiv org abs 1612 08083 where the input x is split in two halves across a dimension axis a and b returning a sigmoid b args x float tensor to perform activation axis dimension across which x be split in half returns x with the glu activation applied with its size halved across the dimension axis
import math

import tensorflow as tf
from packaging import version


def _gelu(x):
    
    x = tf.convert_to_tensor(x)
    cdf = 0.5 * (1.0 + tf.math.erf(x / tf.cast(tf.sqrt(2.0), x.dtype)))

    return x * cdf


def _gelu_new(x):
    
    x = tf.convert_to_tensor(x)
    pi = tf.cast(math.pi, x.dtype)
    coeff = tf.cast(0.044715, x.dtype)
    cdf = 0.5 * (1.0 + tf.tanh(tf.sqrt(2.0 / pi) * (x + coeff * tf.pow(x, 3))))

    return x * cdf


def mish(x):
    x = tf.convert_to_tensor(x)

    return x * tf.tanh(tf.math.softplus(x))


def gelu_fast(x):
    x = tf.convert_to_tensor(x)
    coeff1 = tf.cast(0.044715, x.dtype)
    coeff2 = tf.cast(0.7978845608, x.dtype)

    return 0.5 * x * (1.0 + tf.tanh(x * coeff2 * (1.0 + coeff1 * x * x)))


def quick_gelu(x):
    x = tf.convert_to_tensor(x)
    coeff = tf.cast(1.702, x.dtype)
    return x * tf.math.sigmoid(coeff * x)


def gelu_10(x):
    
    return tf.clip_by_value(_gelu(x), -10, 10)


def glu(x, axis=-1):
    
    a, b = tf.split(x, 2, axis=axis)
    return a * tf.math.sigmoid(b)


if version.parse(tf.version.VERSION) >= version.parse("2.4"):

    def approximate_gelu_wrap(x):
        return tf.keras.activations.gelu(x, approximate=True)

    gelu = tf.keras.activations.gelu
    gelu_new = approximate_gelu_wrap
else:
    gelu = _gelu
    gelu_new = _gelu_new


ACT2FN = {
    "gelu": gelu,
    "gelu_10": gelu_10,
    "gelu_fast": gelu_fast,
    "gelu_new": gelu_new,
    "glu": glu,
    "mish": mish,
    "quick_gelu": quick_gelu,
    "relu": tf.keras.activations.relu,
    "sigmoid": tf.keras.activations.sigmoid,
    "silu": tf.keras.activations.swish,
    "swish": tf.keras.activations.swish,
    "tanh": tf.keras.activations.tanh,
}


def get_tf_activation(activation_string):
    if activation_string in ACT2FN:
        return ACT2FN[activation_string]
    else:
        raise KeyError(f"function {activation_string} not found in ACT2FN mapping {list(ACT2FN.keys())}")