nn

NNX Module 中使用的神经网络层和激活函数。有关更多详细信息，请参阅 NNX 页面。

• 激活函数
  • 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()
• 注意力机制
  • MultiHeadAttention
    • MultiHeadAttention.num_heads
    • MultiHeadAttention.in_features
    • MultiHeadAttention.qkv_features
    • MultiHeadAttention.out_features
    • MultiHeadAttention.dtype
    • MultiHeadAttention.param_dtype
    • MultiHeadAttention.broadcast_dropout
    • MultiHeadAttention.dropout_rate
    • MultiHeadAttention.deterministic
    • MultiHeadAttention.precision
    • MultiHeadAttention.kernel_init
    • MultiHeadAttention.out_kernel_init
    • MultiHeadAttention.bias_init
    • MultiHeadAttention.out_bias_init
    • MultiHeadAttention.use_bias
    • MultiHeadAttention.attention_fn
    • MultiHeadAttention.decode
    • MultiHeadAttention.normalize_qk
    • MultiHeadAttention.rngs
    • MultiHeadAttention.__call__()
    • MultiHeadAttention.init_cache()
  • combine_masks()
  • dot_product_attention()
  • make_attention_mask()
  • make_causal_mask()
• 数据类型
  • canonicalize_dtype()
  • promote_dtype()
• 初始化器
  • 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()
• 线性层
  • Conv
    • Conv.in_features
    • Conv.out_features
    • Conv.kernel_size
    • Conv.strides
    • Conv.padding
    • Conv.input_dilation
    • Conv.kernel_dilation
    • Conv.feature_group_count
    • Conv.use_bias
    • Conv.mask
    • Conv.dtype
    • Conv.param_dtype
    • Conv.precision
    • Conv.kernel_init
    • Conv.bias_init
    • Conv.rngs
    • Conv.__call__()
  • ConvTranspose
    • ConvTranspose.in_features
    • ConvTranspose.out_features
    • ConvTranspose.kernel_size
    • ConvTranspose.strides
    • ConvTranspose.padding
    • ConvTranspose.kernel_dilation
    • ConvTranspose.use_bias
    • ConvTranspose.mask
    • ConvTranspose.dtype
    • ConvTranspose.param_dtype
    • ConvTranspose.precision
    • ConvTranspose.kernel_init
    • ConvTranspose.bias_init
    • ConvTranspose.transpose_kernel
    • ConvTranspose.rngs
    • ConvTranspose.__call__()
  • Embed
    • Embed.num_embeddings
    • Embed.features
    • Embed.dtype
    • Embed.param_dtype
    • Embed.embedding_init
    • Embed.rngs
    • Embed.__call__()
    • Embed.attend()
  • 线性层
    • Linear.in_features
    • Linear.out_features
    • Linear.use_bias
    • Linear.dtype
    • Linear.param_dtype
    • Linear.precision
    • Linear.kernel_init
    • Linear.bias_init
    • Linear.dot_general
    • Linear.rngs
    • Linear.__call__()
  • LinearGeneral
    • LinearGeneral.in_features
    • LinearGeneral.out_features
    • LinearGeneral.axis
    • LinearGeneral.batch_axis
    • LinearGeneral.use_bias
    • LinearGeneral.dtype
    • LinearGeneral.param_dtype
    • LinearGeneral.kernel_init
    • LinearGeneral.bias_init
    • LinearGeneral.precision
    • LinearGeneral.rngs
    • LinearGeneral.__call__()
  • Einsum
    • Einsum.einsum_str
    • Einsum.kernel_shape
    • Einsum.bias_shape
    • Einsum.dtype
    • Einsum.param_dtype
    • Einsum.precision
    • Einsum.kernel_init
    • Einsum.bias_init
    • Einsum.rngs
    • Einsum.__call__()
• LoRA
  • LoRA
    • LoRA.in_features
    • LoRA.lora_rank
    • LoRA.out_features
    • LoRA.base_module
    • LoRA.dtype
    • LoRA.param_dtype
    • LoRA.precision
    • LoRA.kernel_init
    • LoRA.lora_param_type
    • LoRA.__call__()
  • LoRALinear
    • LoRALinear.in_features
    • LoRALinear.out_features
    • LoRALinear.lora_rank
    • LoRALinear.base_module
    • LoRALinear.dtype
    • LoRALinear.param_dtype
    • LoRALinear.precision
    • LoRALinear.kernel_init
    • LoRALinear.lora_param_type
    • LoRALinear.__call__()
• 归一化
  • BatchNorm
    • BatchNorm.num_features
    • BatchNorm.use_running_average
    • BatchNorm.axis
    • BatchNorm.momentum
    • BatchNorm.epsilon
    • BatchNorm.dtype
    • BatchNorm.param_dtype
    • BatchNorm.use_bias
    • BatchNorm.use_scale
    • BatchNorm.bias_init
    • BatchNorm.scale_init
    • BatchNorm.axis_name
    • BatchNorm.axis_index_groups
    • BatchNorm.use_fast_variance
    • BatchNorm.rngs
    • BatchNorm.__call__()
  • LayerNorm
    • LayerNorm.num_features
    • LayerNorm.epsilon
    • LayerNorm.dtype
    • LayerNorm.param_dtype
    • LayerNorm.use_bias
    • LayerNorm.use_scale
    • LayerNorm.bias_init
    • LayerNorm.scale_init
    • LayerNorm.reduction_axes
    • LayerNorm.feature_axes
    • LayerNorm.axis_name
    • LayerNorm.axis_index_groups
    • LayerNorm.use_fast_variance
    • LayerNorm.rngs
    • LayerNorm.__call__()
  • RMSNorm
    • RMSNorm.num_features
    • RMSNorm.epsilon
    • RMSNorm.dtype
    • RMSNorm.param_dtype
    • RMSNorm.use_scale
    • RMSNorm.scale_init
    • RMSNorm.reduction_axes
    • RMSNorm.feature_axes
    • RMSNorm.axis_name
    • RMSNorm.axis_index_groups
    • RMSNorm.use_fast_variance
    • RMSNorm.rngs
    • RMSNorm.__call__()
  • GroupNorm
    • GroupNorm.num_features
    • GroupNorm.num_groups
    • GroupNorm.group_size
    • GroupNorm.epsilon
    • GroupNorm.dtype
    • GroupNorm.param_dtype
    • GroupNorm.use_bias
    • GroupNorm.use_scale
    • GroupNorm.bias_init
    • GroupNorm.scale_init
    • GroupNorm.reduction_axes
    • GroupNorm.axis_name
    • GroupNorm.axis_index_groups
    • GroupNorm.use_fast_variance
    • GroupNorm.rngs
    • GroupNorm.__call__()
• 循环层
  • LSTMCell
    • LSTMCell.__call__()
    • LSTMCell.initialize_carry()
  • OptimizedLSTMCell
    • OptimizedLSTMCell.gate_fn
    • OptimizedLSTMCell.activation_fn
    • OptimizedLSTMCell.kernel_init
    • OptimizedLSTMCell.recurrent_kernel_init
    • OptimizedLSTMCell.bias_init
    • OptimizedLSTMCell.dtype
    • OptimizedLSTMCell.param_dtype
    • OptimizedLSTMCell.__call__()
    • OptimizedLSTMCell.initialize_carry()
  • SimpleCell
    • SimpleCell.__call__()
    • SimpleCell.initialize_carry()
  • GRUCell
    • GRUCell.in_features
    • GRUCell.hidden_features
    • GRUCell.gate_fn
    • GRUCell.activation_fn
    • GRUCell.kernel_init
    • GRUCell.recurrent_kernel_init
    • GRUCell.bias_init
    • GRUCell.dtype
    • GRUCell.param_dtype
    • GRUCell.__call__()
    • GRUCell.initialize_carry()
  • RNN
    • RNN.__call__()
  • Bidirectional
    • Bidirectional.__call__()
  • flip_sequences()
• 随机层
  • Dropout
    • Dropout.rate
    • Dropout.broadcast_dims
    • Dropout.deterministic
    • Dropout.rng_collection
    • Dropout.rngs