Docs updates (#198)

Reorganize NN docs + a few other tidbits.

docs/src/python/nn/functions.rst

@@ -0,0 +1,23 @@
+.. _nn_functions:
+
+.. currentmodule:: mlx.nn
+
+Functions
+---------
+
+Layers without parameters (e.g. activation functions) are also provided as
+simple functions.
+
+.. autosummary::
+   :toctree: _autosummary_functions
+   :template: nn-module-template.rst
+
+   gelu
+   gelu_approx
+   gelu_fast_approx
+   relu
+   prelu
+   silu
+   step
+   selu
+   mish

docs/src/python/nn/layers.rst

@@ -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

docs/src/python/nn/losses.rst

@@ -0,0 +1,17 @@
+.. _losses:
+
+.. currentmodule:: mlx.nn.losses
+
+Loss Functions
+--------------
+
+.. autosummary::
+   :toctree: _autosummary_functions
+   :template: nn-module-template.rst
+
+   cross_entropy
+   binary_cross_entropy
+   l1_loss
+   mse_loss
+   nll_loss
+   kl_div_loss

docs/src/python/nn/module.rst

@@ -1,7 +0,0 @@
-mlx.nn.Module
-=============
-
-.. currentmodule:: mlx.nn
-
-.. autoclass:: Module
-   :members: