flax.linen# Linen 是 Flax 模块系统。在Linen README中阅读更多关于我们的设计目标的信息。 模块 模块 apply() init() init_with_output() intercept_methods() share_scope() 初始化/应用 apply() init() init_with_output() 层 线性模块 池化 归一化 组合器 随机 注意力 循环 BatchApply 激活函数 PReLU celu() elu() gelu() glu() hard_sigmoid() hard_silu() hard_swish() hard_tanh() leaky_relu() log_sigmoid() log_softmax() logsumexp() one_hot() relu() selu() sigmoid() silu() soft_sign() softmax() softplus() standardize() swish() tanh() 初始化器 constant() delta_orthogonal() glorot_normal() glorot_uniform() he_normal() he_uniform() kaiming_normal() kaiming_uniform() lecun_normal() lecun_uniform() normal() truncated_normal() ones() ones_init() orthogonal() uniform() variance_scaling() xavier_normal() xavier_uniform() zeros() zeros_init() 转换 vmap() scan() jit() remat() remat_scan() map_variables() jvp() vjp() custom_vjp() while_loop() cond() switch() 检查 tabulate() 变量字典 变量 SPMD Partitioned() with_partitioning() get_partition_spec() get_sharding() LogicallyPartitioned() logical_axis_rules() set_logical_axis_rules() get_logical_axis_rules() logical_to_mesh_axes() logical_to_mesh() logical_to_mesh_sharding() with_logical_constraint() with_logical_partitioning() 装饰器 compact() nowrap() 性能分析 enable_named_call() disable_named_call() override_named_call()