.. currentmodule:: mlx.nn
.. autosummary:: :toctree: _autosummary :template: nn-module-template.rst ALiBi AllToShardedLinear AvgPool1d AvgPool2d AvgPool3d BatchNorm CELU Conv1d Conv2d Conv3d ConvTranspose1d ConvTranspose2d ConvTranspose3d Dropout Dropout2d Dropout3d Embedding ELU GELU GLU GroupNorm GRU HardShrink HardTanh Hardswish InstanceNorm LayerNorm LeakyReLU Linear LogSigmoid LogSoftmax LSTM MaxPool1d MaxPool2d MaxPool3d Mish MixtureOfExperts MultiHeadAttention PReLU QuantizedAllToShardedLinear QuantizedEmbedding QuantizedLinear QuantizedShardedToAllLinear RMSNorm ReLU ReLU2 ReLU6 RNN RoPE SELU Sequential ShardedToAllLinear Sigmoid SiLU SinusoidalPositionalEncoding Softmin Softshrink Softsign Softmax Softplus Step Tanh Transformer Upsample