nn#
Neural network layers and activation functions used in NNX Module’s.
See the NNX page for more details.
- Activation functions
- Attention
MultiHeadAttentionMultiHeadAttention.num_headsMultiHeadAttention.in_featuresMultiHeadAttention.qkv_featuresMultiHeadAttention.out_featuresMultiHeadAttention.dtypeMultiHeadAttention.param_dtypeMultiHeadAttention.broadcast_dropoutMultiHeadAttention.dropout_rateMultiHeadAttention.deterministicMultiHeadAttention.precisionMultiHeadAttention.kernel_initMultiHeadAttention.out_kernel_initMultiHeadAttention.bias_initMultiHeadAttention.out_bias_initMultiHeadAttention.use_biasMultiHeadAttention.attention_fnMultiHeadAttention.decodeMultiHeadAttention.normalize_qkMultiHeadAttention.rngsMultiHeadAttention.__call__()MultiHeadAttention.init_cache()
combine_masks()dot_product_attention()make_attention_mask()make_causal_mask()
- Initializers
- Linear
ConvConvTransposeConvTranspose.in_featuresConvTranspose.out_featuresConvTranspose.kernel_sizeConvTranspose.stridesConvTranspose.paddingConvTranspose.kernel_dilationConvTranspose.use_biasConvTranspose.maskConvTranspose.dtypeConvTranspose.param_dtypeConvTranspose.precisionConvTranspose.kernel_initConvTranspose.bias_initConvTranspose.transpose_kernelConvTranspose.rngsConvTranspose.__call__()
EmbedLinearLinearGeneralEinsum
- Normalization
BatchNormBatchNorm.num_featuresBatchNorm.use_running_averageBatchNorm.axisBatchNorm.momentumBatchNorm.epsilonBatchNorm.dtypeBatchNorm.param_dtypeBatchNorm.use_biasBatchNorm.use_scaleBatchNorm.bias_initBatchNorm.scale_initBatchNorm.axis_nameBatchNorm.axis_index_groupsBatchNorm.use_fast_varianceBatchNorm.rngsBatchNorm.__call__()
LayerNormLayerNorm.num_featuresLayerNorm.epsilonLayerNorm.dtypeLayerNorm.param_dtypeLayerNorm.use_biasLayerNorm.use_scaleLayerNorm.bias_initLayerNorm.scale_initLayerNorm.reduction_axesLayerNorm.feature_axesLayerNorm.axis_nameLayerNorm.axis_index_groupsLayerNorm.use_fast_varianceLayerNorm.rngsLayerNorm.__call__()
RMSNormGroupNormGroupNorm.num_featuresGroupNorm.num_groupsGroupNorm.group_sizeGroupNorm.epsilonGroupNorm.dtypeGroupNorm.param_dtypeGroupNorm.use_biasGroupNorm.use_scaleGroupNorm.bias_initGroupNorm.scale_initGroupNorm.reduction_axesGroupNorm.axis_nameGroupNorm.axis_index_groupsGroupNorm.use_fast_varianceGroupNorm.rngsGroupNorm.__call__()
- Stochastic