Module checks the weight on load_weights (#337)

* update module to check weights on load, also fix docs and reorganize tests * nits + rebase * a few more docs updates for Module * use manual module file * comment
2025-12-16 01:49:05 +08:00 · 2024-01-02 18:55:42 -08:00
parent 0782a4573a
commit dff4a3833f
6 changed files with 581 additions and 360 deletions
--- a/docs/src/python/nn.rst
+++ b/docs/src/python/nn.rst
@@ -170,14 +170,13 @@ In detail:
  :meth:`mlx.core.value_and_grad`

 .. autosummary::
-   :recursive:
   :toctree: _autosummary

   value_and_grad
-   Module

 .. toctree::

+   nn/module
   nn/layers
   nn/functions
   nn/losses