jax.nn 模块#

神经网络库的常用函数。

激活函数#

relu

整流线性单元激活函数。

relu6

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

sigmoid(x)

Sigmoid 激活函数。

softplus(x)

Softplus 激活函数。

sparse_plus(x)

Sparse plus 函数。

sparse_sigmoid(x)

Sparse sigmoid 激活函数。

soft_sign(x)

Soft-sign 激活函数。

silu(x)

SiLU (又名 swish) 激活函数。

swish(x)

SiLU (又名 swish) 激活函数。

log_sigmoid(x)

Log-sigmoid 激活函数。

leaky_relu(x[, negative_slope])

Leaky rectified linear unit 激活函数。

hard_sigmoid(x)

Hard Sigmoid 激活函数。

hard_silu(x)

Hard SiLU (swish) 激活函数

hard_swish(x)

Hard SiLU (swish) 激活函数

hard_tanh(x)

Hard \(\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)

Identity 激活函数。

其他函数#

softmax(x[, axis, where])

Softmax 函数。

log_softmax(x[, axis, where])

Log-Softmax 函数。

logsumexp()

Log-sum-exp 归约。

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

通过减去 mean 并除以 \(\sqrt{\mathrm{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)

缩放点积通用操作。