| 33 | // Functor used by GeluOp to do the computations. |
| 34 | template <typename Device, typename T> |
| 35 | struct Gelu { |
| 36 | // Computes Gelu activation. |
| 37 | // |
| 38 | // features: any shape. |
| 39 | // approximate: whether to enable approximation. |
| 40 | // activations: same shape as "features". |
| 41 | void operator()(const Device& d, typename TTypes<T>::ConstTensor features, |
| 42 | bool approximate, typename TTypes<T>::Tensor activations) { |
| 43 | const T one = static_cast<T>(1); |
| 44 | const T half = static_cast<T>(0.5); |
| 45 | if (approximate) { |
| 46 | // y = 0.5 * x * (1 + tanh(sqrt(2 / pi) * (x + 0.044715 * x^3))) |
| 47 | activations.device(d) = |
| 48 | half * features * |
| 49 | (one + |
| 50 | (static_cast<T>(internal::kAlpha) * |
| 51 | (features + static_cast<T>(internal::kCoeff) * features.cube())) |
| 52 | .tanh()); |
| 53 | } else { |
| 54 | // y = x * normcdf(x) = 0.5 * x * (1 + erf(x / sqrt(2))) |
| 55 | activations.device(d) = |
| 56 | half * features * |
| 57 | (one + (features * static_cast<T>(internal::kSqrtHalf)).erf()); |
| 58 | } |
| 59 | } |
| 60 | }; |
| 61 | |
| 62 | // Functor used by GeluGradOp to do the computations. |
| 63 | template <typename Device, typename T> |
no outgoing calls