JAX-based activation function, softmax function and cross entropy function

1.tanh激活函数公式

$$\tanh \left(x\right)=\frac{e^x-e^{-x}}{e^x+e^{-x}}$$

2.基于JAX实现tanh激活函数

import jax.numpy as jnp
from jax import random
def tanh(x):
    """tanh function"""
    return (jnp.exp(x)-jnp.exp(-x))/(jnp.exp(x)+jnp.exp(-x))
# Set the pseudorandom number seed
rng=random.PRNGKey(0)
# Standard normal sampling yields the input vector
x=random.normal(rng,shape=(4,1))
# 调用内置tanh函数实现
print(jnp.tanh(x))
# 调用自定义tanh函数实现
print(tanh(x))

3.softmax函数公式

$$s_i=\frac{e^{x_i}}{\sum _{i=0}^j{e}^{x_i}}$$

其中$x_i$表示第$i$real-valued output of a neuron

4.基于JAX实现softmax函数

import jax.numpy as jnp
import jax.nn as nn
def softmax(x,axis=-1):
    """softmax function"""
    unnormalized=jnp.exp(x)
    return unnormalized/jnp.sum(unnormalized)
# 定义数组
arr=jnp.arange(-2,4)
# 调用自定义softmax函数
print(softmax(arr))
# 调用jax自带softmax函数
print(nn.softmax(arr))   

5.Cross-entropy function formula

$$H\left(p,q\right)=-\sum _{i=1}^{n}p\left(x_i\right)\log \left(q\left(x_i\right)\right)$$

其中$p(x)$represents the true probability distribution,$q(x)$represents the predicted probability distribution

6.基于JAXImplement the cross-entropy function

import jax.numpy as jnp
def cross_entropy(y_true,y_pred,eps=1e-7):
    """cross entropy function :param y_true:真实标签 :param y_pred:Neural network predicts labels :param eps:Default minimal positive number,The logarithm is guaranteed to be true0,增强logFunction Numerical Stability :return:交叉熵,保留到小数点后4位 """
    y_true=jnp.array(y_true)
    y_pred=jnp.array(y_pred)
    res=-jnp.sum(y_true*jnp.log(y_pred+eps),axis=-1)
    return jnp.round(res,4)
# 预测概率分布
y_pred=[0.1,0.05,0.6,0.0,0.05,0.1,0.0,0.1,0.0,0.0]
# 真实概率分布
y_true=[0,0,1,0,0,0,0,0,0,0]
# 交叉熵为0.5108
print(cross_entropy(y_true,y_pred))