docs/src/python/nn/layers.rst

.. _layers:

.. currentmodule:: mlx.nn

Layers
------

.. autosummary::
   :toctree: _autosummary
   :template: nn-module-template.rst

   ALiBi
   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
   MultiHeadAttention
   PReLU
   QuantizedEmbedding
   QuantizedLinear
   RMSNorm
   ReLU
   ReLU2
   ReLU6
   RNN
   RoPE
   SELU
   Sequential
   Sigmoid
   SiLU
   SinusoidalPositionalEncoding
   Softmin
   Softshrink
   Softsign
   Softmax
   Softplus
   Step
   Tanh
   Transformer
   Upsample
Docs updates (#198) Reorganize NN docs + a few other tidbits. 2023-12-17 13:20:55 -08:00			`.. _layers:`

			`.. currentmodule:: mlx.nn`

			`Layers`
			`------`

			`.. autosummary::`
			`:toctree: _autosummary`
			`:template: nn-module-template.rst`

Sort some APIs docs by names (a-z) (#472) 2024-01-17 11:37:50 +08:00			`ALiBi`
Pooling layers (#357) Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com> Co-authored-by: Awni Hannun <awni@apple.com> 2024-02-13 07:08:13 +01:00			`AvgPool1d`
			`AvgPool2d`
Doc fix (#1615) 2024-11-22 11:12:25 -08:00			`AvgPool3d`
Sort some APIs docs by names (a-z) (#472) 2024-01-17 11:37:50 +08:00			`BatchNorm`
More fixes for arrays with large sizes (#1405) * compile works for big arrays when contiguous * style * nits in docs * a bunch more stuff * update jit * update jit * use constant for shapes and strides and remove elem_to_loc overload * use kernel instantiation * docs nits * update binary and ternary * comments 2024-09-17 12:46:31 -07:00			`CELU`
Docs updates (#198) Reorganize NN docs + a few other tidbits. 2023-12-17 13:20:55 -08:00			`Conv1d`
			`Conv2d`
Conv3d (#993) * added conv3d added conv3d implemented explicit_gemm_conv_ND_cpu and bounds checks for slow_conv_3D * incorporated reviewer comments * fixed test * reduced tensor shapes in test for conv3d * Reviewer suggestion Co-authored-by: Awni Hannun <awni.hannun@gmail.com> Reviewer suggestion Co-authored-by: Awni Hannun <awni.hannun@gmail.com> Reviewer suggestion Co-authored-by: Awni Hannun <awni.hannun@gmail.com> Reviewer suggestion 2024-05-11 15:15:02 +02:00			`Conv3d`
Transposed Convolution (#1245) * initial implementation for conv_transpose ran pre-commit implemented conv_transpose updated conv_general docstring updated conv_general docstring updated code comments removed commented run_conv_checks updated acknowledgments added missing entry to ops.rst added op to nn.layers resolved merge conflicts * removed ConvolutionTranspose primitive as suggested by reviewer removed ConvolutionTranspose primitive as suggested by reviewer * remove transpose flag, add another test --------- Co-authored-by: Awni Hannun <awni@apple.com> 2024-09-07 04:52:38 +02:00			`ConvTranspose1d`
			`ConvTranspose2d`
			`ConvTranspose3d`
Add dropout2d (#250) 2023-12-22 17:02:29 +01:00			`Dropout`
			`Dropout2d`
feat: Add Dropout3d layer to nn.layers (#313) * feat: Add Dropout3d layer to nn.layers * acknowledgement * Add dropout tests to test_nn.py * run pre-commit * Add activation functions and dropout3d ops * Add dropout tests for bfloat16 and float16 2024-01-01 02:01:21 +04:00			`Dropout3d`
Sort some APIs docs by names (a-z) (#472) 2024-01-17 11:37:50 +08:00			`Embedding`
More fixes for arrays with large sizes (#1405) * compile works for big arrays when contiguous * style * nits in docs * a bunch more stuff * update jit * update jit * use constant for shapes and strides and remove elem_to_loc overload * use kernel instantiation * docs nits * update binary and ternary * comments 2024-09-17 12:46:31 -07:00			`ELU`
Sort some APIs docs by names (a-z) (#472) 2024-01-17 11:37:50 +08:00			`GELU`
Add softmin, hardshrink, hardtanh (#1180) --------- Co-authored-by: Nikhil Mehta <nikmehta@tesla.com> 2024-06-04 15:48:18 -07:00			`GLU`
Sort some APIs docs by names (a-z) (#472) 2024-01-17 11:37:50 +08:00			`GroupNorm`
Implement RNN, GRU, LSTM (#268) * RNN base implementation * Address comments+format * nits in docs * add tests for prb * fix test * add a couple tests --------- Co-authored-by: Awni Hannun <awni@apple.com> 2024-03-12 05:14:44 +01:00			`GRU`
Add softmin, hardshrink, hardtanh (#1180) --------- Co-authored-by: Nikhil Mehta <nikmehta@tesla.com> 2024-06-04 15:48:18 -07:00			`HardShrink`
			`HardTanh`
			`Hardswish`
Sort some APIs docs by names (a-z) (#472) 2024-01-17 11:37:50 +08:00			`InstanceNorm`
			`LayerNorm`
Add softmin, hardshrink, hardtanh (#1180) --------- Co-authored-by: Nikhil Mehta <nikmehta@tesla.com> 2024-06-04 15:48:18 -07:00			`LeakyReLU`
Sort some APIs docs by names (a-z) (#472) 2024-01-17 11:37:50 +08:00			`Linear`
More fixes for arrays with large sizes (#1405) * compile works for big arrays when contiguous * style * nits in docs * a bunch more stuff * update jit * update jit * use constant for shapes and strides and remove elem_to_loc overload * use kernel instantiation * docs nits * update binary and ternary * comments 2024-09-17 12:46:31 -07:00			`LogSigmoid`
			`LogSoftmax`
Implement RNN, GRU, LSTM (#268) * RNN base implementation * Address comments+format * nits in docs * add tests for prb * fix test * add a couple tests --------- Co-authored-by: Awni Hannun <awni@apple.com> 2024-03-12 05:14:44 +01:00			`LSTM`
Pooling layers (#357) Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com> Co-authored-by: Awni Hannun <awni@apple.com> 2024-02-13 07:08:13 +01:00			`MaxPool1d`
			`MaxPool2d`
Doc fix (#1615) 2024-11-22 11:12:25 -08:00			`MaxPool3d`
Sort some APIs docs by names (a-z) (#472) 2024-01-17 11:37:50 +08:00			`Mish`
Transformer fix (#167) * add transformer with dropout, fix transformer ffm, layernorm order * precommit changes * precommit changes * add docstring, activation, norm_first * run precommit * run precommit * add doctstring * precommit * style nits in docs --------- Co-authored-by: junwoo-yun <junwoo.yun@bagelcode.com> Co-authored-by: Awni Hannun <awni@apple.com> 2023-12-28 00:48:36 +08:00			`MultiHeadAttention`
Sort some APIs docs by names (a-z) (#472) 2024-01-17 11:37:50 +08:00			`PReLU`
Quantize embedding (#994) * quantize embedding * rename as_linear + comment * consistency in docs * fix test 2024-04-15 16:42:10 -07:00			`QuantizedEmbedding`
Sort some APIs docs by names (a-z) (#472) 2024-01-17 11:37:50 +08:00			`QuantizedLinear`
			`RMSNorm`
			`ReLU`
Adding Relu2 (#2582) * in. com. * upd. ackn. * update __init__ * nits * nits + format * used mx.maximum(x, 0) instead of calling the function and moves relu6 under relu2 to make it nicer * same with _make_activation_module * Update python/mlx/nn/layers/activations.py upd Co-authored-by: Awni Hannun <awni.hannun@gmail.com> * update funct.rst * upd. layers.rst --------- Co-authored-by: Awni Hannun <awni.hannun@gmail.com> 2025-09-10 16:24:30 +02:00			`ReLU2`
Add softmin, hardshrink, hardtanh (#1180) --------- Co-authored-by: Nikhil Mehta <nikmehta@tesla.com> 2024-06-04 15:48:18 -07:00			`ReLU6`
Implement RNN, GRU, LSTM (#268) * RNN base implementation * Address comments+format * nits in docs * add tests for prb * fix test * add a couple tests --------- Co-authored-by: Awni Hannun <awni@apple.com> 2024-03-12 05:14:44 +01:00			`RNN`
Transformer fix (#167) * add transformer with dropout, fix transformer ffm, layernorm order * precommit changes * precommit changes * add docstring, activation, norm_first * run precommit * run precommit * add doctstring * precommit * style nits in docs --------- Co-authored-by: junwoo-yun <junwoo.yun@bagelcode.com> Co-authored-by: Awni Hannun <awni@apple.com> 2023-12-28 00:48:36 +08:00			`RoPE`
Sort some APIs docs by names (a-z) (#472) 2024-01-17 11:37:50 +08:00			`SELU`
			`Sequential`
More fixes for arrays with large sizes (#1405) * compile works for big arrays when contiguous * style * nits in docs * a bunch more stuff * update jit * update jit * use constant for shapes and strides and remove elem_to_loc overload * use kernel instantiation * docs nits * update binary and ternary * comments 2024-09-17 12:46:31 -07:00			`Sigmoid`
Sort some APIs docs by names (a-z) (#472) 2024-01-17 11:37:50 +08:00			`SiLU`
Transformer fix (#167) * add transformer with dropout, fix transformer ffm, layernorm order * precommit changes * precommit changes * add docstring, activation, norm_first * run precommit * run precommit * add doctstring * precommit * style nits in docs --------- Co-authored-by: junwoo-yun <junwoo.yun@bagelcode.com> Co-authored-by: Awni Hannun <awni@apple.com> 2023-12-28 00:48:36 +08:00			`SinusoidalPositionalEncoding`
Add softmin, hardshrink, hardtanh (#1180) --------- Co-authored-by: Nikhil Mehta <nikmehta@tesla.com> 2024-06-04 15:48:18 -07:00			`Softmin`
Softshrink mapping + op (#552) * Added Softshrink mapping + op * formatting * docs + nits in docstring --------- Co-authored-by: Awni Hannun <awni@apple.com> 2024-01-30 15:56:28 -05:00			`Softshrink`
Add softmin, hardshrink, hardtanh (#1180) --------- Co-authored-by: Nikhil Mehta <nikmehta@tesla.com> 2024-06-04 15:48:18 -07:00			`Softsign`
			`Softmax`
			`Softplus`
Sort some APIs docs by names (a-z) (#472) 2024-01-17 11:37:50 +08:00			`Step`
Add softmin, hardshrink, hardtanh (#1180) --------- Co-authored-by: Nikhil Mehta <nikmehta@tesla.com> 2024-06-04 15:48:18 -07:00			`Tanh`
Sort some APIs docs by names (a-z) (#472) 2024-01-17 11:37:50 +08:00			`Transformer`
Quantize embedding (#994) * quantize embedding * rename as_linear + comment * consistency in docs * fix test 2024-04-15 16:42:10 -07:00			`Upsample`