jax.nn 模块#

神经网络库的常用函数。

激活函数#

relu

整流线性单元激活函数。

relu6

整流线性单元 6 激活函数。

sigmoid(x)

Sigmoid 激活函数。

softplus(x)

Softplus 激活函数。

sparse_plus(x)

稀疏加法函数。

sparse_sigmoid(x)

稀疏 Sigmoid 激活函数。

soft_sign(x)

Soft-sign 激活函数。

silu(x)

SiLU（又称 swish）激活函数。

swish(x)

SiLU（又称 swish）激活函数。

log_sigmoid(x)

Log-sigmoid 激活函数。

leaky_relu(x[, negative_slope])

泄漏整流线性单元激活函数。

hard_sigmoid(x)

硬 Sigmoid 激活函数。

hard_silu(x)

硬 SiLU (swish) 激活函数

hard_swish(x)

硬 SiLU (swish) 激活函数

hard_tanh(x)

硬 \(\mathrm{tanh}\) 激活函数。

elu(x[, alpha])

指数线性单元激活函数。

celu(x[, alpha])

连续可微指数线性单元激活。

selu(x)

缩放指数线性单元激活。

gelu(x[, approximate])

高斯误差线性单元激活函数。

glu(x[, axis])

门控线性单元激活函数。

squareplus(x[, b])

Squareplus 激活函数。

mish(x)

Mish 激活函数。

identity(x)

恒等激活函数。

其他函数#

softmax(x[, axis, where])

Softmax 函数。

log_softmax(x[, axis, where])

Log-Softmax 函数。

logsumexp()

对数和指数缩减。

standardize(x[, axis, mean, variance, ...])

将输入标准化为零均值和单位方差。

one_hot(x, num_classes, *[, dtype, axis])

对给定索引进行独热编码。

dot_product_attention(query, key, value[, ...])

缩放点积注意力函数。

scaled_matmul(lhs, rhs, lhs_scales, rhs_scales)

缩放矩阵乘法函数。

get_scaled_dot_general_config(mode[, ...])

获取 scaled_dot_general 的量化配置。

scaled_dot_general(lhs, rhs, dimension_numbers)

缩放点积通用操作。