Docs updates (#198)

Reorganize NN docs + a few other tidbits.
2025-12-16 01:49:05 +08:00 · 2023-12-17 13:20:55 -08:00
parent 90d04072b7
commit ee0c2835c5
9 changed files with 160 additions and 74 deletions
--- a/docs/src/python/nn/losses.rst
+++ b/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