Docs updates (#198)

Reorganize NN docs + a few other tidbits.
2023-12-17 13:20:55 -08:00
parent 90d04072b7
commit ee0c2835c5
9 changed files with 160 additions and 74 deletions
@@ -0,0 +1,28 @@
+.. _layers:
+
+.. currentmodule:: mlx.nn
+
+Layers
+------
+
+.. autosummary::
+   :toctree: _autosummary
+   :template: nn-module-template.rst
+
+   Embedding
+   ReLU
+   PReLU
+   GELU
+   SiLU
+   Step
+   SELU
+   Mish
+   Linear
+   Conv1d
+   Conv2d
+   LayerNorm
+   RMSNorm
+   GroupNorm
+   RoPE
+   MultiHeadAttention
+   Sequential