Jax calculation Selu function

1.SeLU(scaled exponential linear units) Activate function calculation formula

$$\text{selu}(x)=\lambda \{\begin{array}{ll}x& \text{ if }x>0\\ \alpha e^x-\alpha & \text{ if }x⩽0.\end{array}$$

among $\lambda =1.0507009873554804934193349852946$,$\alpha =\mathrm{1.6732632423543772848170429916717.}$

2.JAX Code implementation

#!/usr/bin/python3
# -*- coding: utf-8 -*-
""" @Time : 2022/7/20 13:51 @Author : Albert Darren @Contact : [email protected] @File : Program1.1.py @Version : Version 1.0.0 @Description : TODO  utilize jax Calculation selu function , See P12 @Created By : PyCharm """
import jax.numpy as jnp  #  Import numpy Calculation package 
from jax import random  #  Import random Random number packets 


def selu(x, alpha=1.6732632423543772848170429916717, lmbda=1.0507009873554804934193349852946):
    """  Realization selu Activation function  :param x:  Input tensor  :param alpha:  Predefined parameters alpha :param lmbda:  Predefined parameters lambda, The variable name here is intentionally misspelled , Avoid with keywords lambda name conflict  :return: selu Function value  """
    return lmbda * jnp.where(x > 0, x, alpha * (jnp.exp(x) - 1))


#  Generate a fixed number 17 As key
key = random.PRNGKey(17)
#  Randomly generate a size of [1,5] Matrix 
x = random.normal(key, (5,))
print(selu(x))
# [-1.2497659 0.4546819 1.5760192 -0.81573856 0.27510932]

3. reference

• Self-Normalizing Neural Networks
• Google JAX Deep learning starts from scratch