docs update

docs/build/html/_sources/python/nn/_autosummary/mlx.nn.ALiBi.rst

@@ -0,0 +1,8 @@
+mlx.nn.ALiBi
+============
+
+.. currentmodule:: mlx.nn
+
+.. autoclass:: ALiBi
+
+   

docs/build/html/_sources/python/nn/_autosummary/mlx.nn.BatchNorm.rst

@@ -0,0 +1,8 @@
+mlx.nn.BatchNorm
+================
+
+.. currentmodule:: mlx.nn
+
+.. autoclass:: BatchNorm
+
+   

docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Dropout.rst

@@ -0,0 +1,8 @@
+mlx.nn.Dropout
+==============
+
+.. currentmodule:: mlx.nn
+
+.. autoclass:: Dropout
+
+   

docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Dropout2d.rst

@@ -0,0 +1,8 @@
+mlx.nn.Dropout2d
+================
+
+.. currentmodule:: mlx.nn
+
+.. autoclass:: Dropout2d
+
+   

docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Dropout3d.rst

@@ -0,0 +1,8 @@
+mlx.nn.Dropout3d
+================
+
+.. currentmodule:: mlx.nn
+
+.. autoclass:: Dropout3d
+
+   

docs/build/html/_sources/python/nn/_autosummary/mlx.nn.InstanceNorm.rst

@@ -0,0 +1,8 @@
+mlx.nn.InstanceNorm
+===================
+
+.. currentmodule:: mlx.nn
+
+.. autoclass:: InstanceNorm
+
+   

docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Module.apply.rst

@@ -0,0 +1,6 @@
+mlx.nn.Module.apply
+===================
+
+.. currentmodule:: mlx.nn
+
+.. automethod:: Module.apply

docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Module.apply_to_modules.rst

@@ -0,0 +1,6 @@
+mlx.nn.Module.apply\_to\_modules
+================================
+
+.. currentmodule:: mlx.nn
+
+.. automethod:: Module.apply_to_modules

docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Module.children.rst

@@ -0,0 +1,6 @@
+mlx.nn.Module.children
+======================
+
+.. currentmodule:: mlx.nn
+
+.. automethod:: Module.children

docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Module.eval.rst

@@ -0,0 +1,6 @@
+mlx.nn.Module.eval
+==================
+
+.. currentmodule:: mlx.nn
+
+.. automethod:: Module.eval

docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Module.filter_and_map.rst

@@ -0,0 +1,6 @@
+mlx.nn.Module.filter\_and\_map
+==============================
+
+.. currentmodule:: mlx.nn
+
+.. automethod:: Module.filter_and_map

docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Module.freeze.rst

@@ -0,0 +1,6 @@
+mlx.nn.Module.freeze
+====================
+
+.. currentmodule:: mlx.nn
+
+.. automethod:: Module.freeze

docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Module.leaf_modules.rst

@@ -0,0 +1,6 @@
+mlx.nn.Module.leaf\_modules
+===========================
+
+.. currentmodule:: mlx.nn
+
+.. automethod:: Module.leaf_modules

docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Module.load_weights.rst

@@ -0,0 +1,6 @@
+mlx.nn.Module.load\_weights
+===========================
+
+.. currentmodule:: mlx.nn
+
+.. automethod:: Module.load_weights

docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Module.modules.rst

@@ -0,0 +1,6 @@
+mlx.nn.Module.modules
+=====================
+
+.. currentmodule:: mlx.nn
+
+.. automethod:: Module.modules

docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Module.named_modules.rst

@@ -0,0 +1,6 @@
+mlx.nn.Module.named\_modules
+============================
+
+.. currentmodule:: mlx.nn
+
+.. automethod:: Module.named_modules

docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Module.parameters.rst

@@ -0,0 +1,6 @@
+mlx.nn.Module.parameters
+========================
+
+.. currentmodule:: mlx.nn
+
+.. automethod:: Module.parameters

docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Module.save_weights.rst

@@ -0,0 +1,6 @@
+mlx.nn.Module.save\_weights
+===========================
+
+.. currentmodule:: mlx.nn
+
+.. automethod:: Module.save_weights

docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Module.train.rst

@@ -0,0 +1,6 @@
+mlx.nn.Module.train
+===================
+
+.. currentmodule:: mlx.nn
+
+.. automethod:: Module.train

docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Module.trainable_parameters.rst

@@ -0,0 +1,6 @@
+mlx.nn.Module.trainable\_parameters
+===================================
+
+.. currentmodule:: mlx.nn
+
+.. automethod:: Module.trainable_parameters

docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Module.training.rst

@@ -0,0 +1,6 @@
+mlx.nn.Module.training
+======================
+
+.. currentmodule:: mlx.nn
+
+.. autoproperty:: Module.training

docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Module.unfreeze.rst

@@ -0,0 +1,6 @@
+mlx.nn.Module.unfreeze
+======================
+
+.. currentmodule:: mlx.nn
+
+.. automethod:: Module.unfreeze

docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Module.update.rst

@@ -0,0 +1,6 @@
+mlx.nn.Module.update
+====================
+
+.. currentmodule:: mlx.nn
+
+.. automethod:: Module.update

docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Module.update_modules.rst

@@ -0,0 +1,6 @@
+mlx.nn.Module.update\_modules
+=============================
+
+.. currentmodule:: mlx.nn
+
+.. automethod:: Module.update_modules

docs/build/html/_sources/python/nn/_autosummary/mlx.nn.SinusoidalPositionalEncoding.rst

@@ -0,0 +1,8 @@
+mlx.nn.SinusoidalPositionalEncoding
+===================================
+
+.. currentmodule:: mlx.nn
+
+.. autoclass:: SinusoidalPositionalEncoding
+
+   

docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Transformer.rst

@@ -0,0 +1,8 @@
+mlx.nn.Transformer
+==================
+
+.. currentmodule:: mlx.nn
+
+.. autoclass:: Transformer
+
+   

docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.losses.hinge_loss.rst

@@ -0,0 +1,8 @@
+mlx.nn.losses.hinge\_loss
+=========================
+
+.. currentmodule:: mlx.nn.losses
+
+.. autoclass:: hinge_loss
+
+   

docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.losses.huber_loss.rst

@@ -0,0 +1,8 @@
+mlx.nn.losses.huber\_loss
+=========================
+
+.. currentmodule:: mlx.nn.losses
+
+.. autoclass:: huber_loss
+
+   

docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.losses.log_cosh_loss.rst

@@ -0,0 +1,8 @@
+mlx.nn.losses.log\_cosh\_loss
+=============================
+
+.. currentmodule:: mlx.nn.losses
+
+.. autoclass:: log_cosh_loss
+
+   

docs/build/html/_sources/python/nn/layers.rst

@@ -9,7 +9,7 @@ Layers
    :toctree: _autosummary
    :template: nn-module-template.rst
 
-   Embedding
+   Sequential
    ReLU
    PReLU
    GELU
@@ -17,13 +17,21 @@ Layers
    Step
    SELU
    Mish
+   Embedding
    Linear
+   QuantizedLinear
    Conv1d
    Conv2d
+   BatchNorm
    LayerNorm
    RMSNorm
    GroupNorm
-   RoPE
+   InstanceNorm
+   Dropout
+   Dropout2d
+   Dropout3d
+   Transformer
    MultiHeadAttention
-   Sequential
-   QuantizedLinear
+   ALiBi
+   RoPE
+   SinusoidalPositionalEncoding

docs/build/html/_sources/python/nn/losses.rst

@@ -16,4 +16,7 @@ Loss Functions
    mse_loss
    nll_loss
    smooth_l1_loss
-   triplet_loss
+   triplet_loss
+   hinge_loss
+   huber_loss
+   log_cosh_loss

docs/build/html/_sources/python/nn/module.rst

@@ -0,0 +1,36 @@
+Module
+======
+
+.. currentmodule:: mlx.nn
+
+.. autoclass:: Module
+
+   .. rubric:: Attributes
+
+   .. autosummary::
+      :toctree: _autosummary
+   
+      Module.training
+   
+   .. rubric:: Methods
+
+   .. autosummary::
+      :toctree: _autosummary
+   
+      Module.apply
+      Module.apply_to_modules
+      Module.children
+      Module.eval
+      Module.filter_and_map
+      Module.freeze
+      Module.leaf_modules
+      Module.load_weights
+      Module.modules
+      Module.named_modules
+      Module.parameters
+      Module.save_weights
+      Module.train
+      Module.trainable_parameters
+      Module.unfreeze
+      Module.update
+      Module.update_modules