docs update

2025-09-19 10:48:09 +08:00 · 2024-04-26 08:24:09 -07:00
parent ba26fe5be5
commit 61c03b86a4
502 changed files with 14665 additions and 15103 deletions
--- a/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.ALiBi.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.ALiBi.rst
@@ -3,4 +3,16 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: ALiBi
+.. autoclass:: ALiBi
+
+   
+
+   
+   .. rubric:: Methods
+
+   .. autosummary::
+   
+         ~ALiBi.create_alibi_matrix
+         ~ALiBi.create_alibi_slope
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.AvgPool1d.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.AvgPool1d.rst
@@ -3,4 +3,14 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: AvgPool1d
+.. autoclass:: AvgPool1d
+
+   
+
+   
+   .. rubric:: Methods
+
+   .. autosummary::
+   
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.AvgPool2d.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.AvgPool2d.rst
@@ -3,4 +3,14 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: AvgPool2d
+.. autoclass:: AvgPool2d
+
+   
+
+   
+   .. rubric:: Methods
+
+   .. autosummary::
+   
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.BatchNorm.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.BatchNorm.rst
@@ -3,4 +3,15 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: BatchNorm
+.. autoclass:: BatchNorm
+
+   
+
+   
+   .. rubric:: Methods
+
+   .. autosummary::
+   
+         ~BatchNorm.unfreeze
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Conv1d.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Conv1d.rst
@@ -3,4 +3,14 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: Conv1d
+.. autoclass:: Conv1d
+
+   
+
+   
+   .. rubric:: Methods
+
+   .. autosummary::
+   
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Conv2d.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Conv2d.rst
@@ -3,4 +3,14 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: Conv2d
+.. autoclass:: Conv2d
+
+   
+
+   
+   .. rubric:: Methods
+
+   .. autosummary::
+   
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Dropout.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Dropout.rst
@@ -3,4 +3,14 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: Dropout
+.. autoclass:: Dropout
+
+   
+
+   
+   .. rubric:: Methods
+
+   .. autosummary::
+   
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Dropout2d.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Dropout2d.rst
@@ -3,4 +3,14 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: Dropout2d
+.. autoclass:: Dropout2d
+
+   
+
+   
+   .. rubric:: Methods
+
+   .. autosummary::
+   
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Dropout3d.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Dropout3d.rst
@@ -3,4 +3,14 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: Dropout3d
+.. autoclass:: Dropout3d
+
+   
+
+   
+   .. rubric:: Methods
+
+   .. autosummary::
+   
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Embedding.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Embedding.rst
@@ -3,4 +3,15 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: Embedding
+.. autoclass:: Embedding
+
+   
+
+   
+   .. rubric:: Methods
+
+   .. autosummary::
+   
+         ~Embedding.as_linear
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.GELU.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.GELU.rst
@@ -3,4 +3,14 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: GELU
+.. autoclass:: GELU
+
+   
+
+   
+   .. rubric:: Methods
+
+   .. autosummary::
+   
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.GRU.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.GRU.rst
@@ -3,4 +3,14 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: GRU
+.. autoclass:: GRU
+
+   
+
+   
+   .. rubric:: Methods
+
+   .. autosummary::
+   
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.GroupNorm.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.GroupNorm.rst
@@ -3,4 +3,14 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: GroupNorm
+.. autoclass:: GroupNorm
+
+   
+
+   
+   .. rubric:: Methods
+
+   .. autosummary::
+   
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.InstanceNorm.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.InstanceNorm.rst
@@ -3,4 +3,14 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: InstanceNorm
+.. autoclass:: InstanceNorm
+
+   
+
+   
+   .. rubric:: Methods
+
+   .. autosummary::
+   
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.LSTM.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.LSTM.rst
@@ -3,4 +3,14 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: LSTM
+.. autoclass:: LSTM
+
+   
+
+   
+   .. rubric:: Methods
+
+   .. autosummary::
+   
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.LayerNorm.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.LayerNorm.rst
@@ -3,4 +3,14 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: LayerNorm
+.. autoclass:: LayerNorm
+
+   
+
+   
+   .. rubric:: Methods
+
+   .. autosummary::
+   
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Linear.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Linear.rst
@@ -3,4 +3,14 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: Linear
+.. autoclass:: Linear
+
+   
+
+   
+   .. rubric:: Methods
+
+   .. autosummary::
+   
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.MaxPool1d.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.MaxPool1d.rst
@@ -3,4 +3,14 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: MaxPool1d
+.. autoclass:: MaxPool1d
+
+   
+
+   
+   .. rubric:: Methods
+
+   .. autosummary::
+   
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.MaxPool2d.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.MaxPool2d.rst
@@ -3,4 +3,14 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: MaxPool2d
+.. autoclass:: MaxPool2d
+
+   
+
+   
+   .. rubric:: Methods
+
+   .. autosummary::
+   
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Mish.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Mish.rst
@@ -3,4 +3,14 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: Mish
+.. autoclass:: Mish
+
+   
+
+   
+   .. rubric:: Methods
+
+   .. autosummary::
+   
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.MultiHeadAttention.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.MultiHeadAttention.rst
@@ -3,4 +3,15 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: MultiHeadAttention
+.. autoclass:: MultiHeadAttention
+
+   
+
+   
+   .. rubric:: Methods
+
+   .. autosummary::
+   
+         ~MultiHeadAttention.create_additive_causal_mask
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.PReLU.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.PReLU.rst
@@ -3,4 +3,14 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: PReLU
+.. autoclass:: PReLU
+
+   
+
+   
+   .. rubric:: Methods
+
+   .. autosummary::
+   
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.QuantizedEmbedding.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.QuantizedEmbedding.rst
@@ -0,0 +1,18 @@
+mlx.nn.QuantizedEmbedding
+=========================
+
+.. currentmodule:: mlx.nn
+
+.. autoclass:: QuantizedEmbedding
+
+   
+
+   
+   .. rubric:: Methods
+
+   .. autosummary::
+   
+         ~QuantizedEmbedding.as_linear
+         ~QuantizedEmbedding.from_embedding
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.QuantizedLinear.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.QuantizedLinear.rst
@@ -3,4 +3,16 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: QuantizedLinear
+.. autoclass:: QuantizedLinear
+
+   
+
+   
+   .. rubric:: Methods
+
+   .. autosummary::
+   
+         ~QuantizedLinear.from_linear
+         ~QuantizedLinear.unfreeze
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.RMSNorm.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.RMSNorm.rst
@@ -3,4 +3,14 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: RMSNorm
+.. autoclass:: RMSNorm
+
+   
+
+   
+   .. rubric:: Methods
+
+   .. autosummary::
+   
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.RNN.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.RNN.rst
@@ -3,4 +3,14 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: RNN
+.. autoclass:: RNN
+
+   
+
+   
+   .. rubric:: Methods
+
+   .. autosummary::
+   
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.ReLU.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.ReLU.rst
@@ -3,4 +3,14 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: ReLU
+.. autoclass:: ReLU
+
+   
+
+   
+   .. rubric:: Methods
+
+   .. autosummary::
+   
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.RoPE.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.RoPE.rst
@@ -3,4 +3,14 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: RoPE
+.. autoclass:: RoPE
+
+   
+
+   
+   .. rubric:: Methods
+
+   .. autosummary::
+   
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.SELU.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.SELU.rst
@@ -3,4 +3,14 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: SELU
+.. autoclass:: SELU
+
+   
+
+   
+   .. rubric:: Methods
+
+   .. autosummary::
+   
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Sequential.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Sequential.rst
@@ -3,4 +3,14 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: Sequential
+.. autoclass:: Sequential
+
+   
+
+   
+   .. rubric:: Methods
+
+   .. autosummary::
+   
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.SiLU.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.SiLU.rst
@@ -3,4 +3,14 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: SiLU
+.. autoclass:: SiLU
+
+   
+
+   
+   .. rubric:: Methods
+
+   .. autosummary::
+   
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.SinusoidalPositionalEncoding.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.SinusoidalPositionalEncoding.rst
@@ -3,4 +3,14 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: SinusoidalPositionalEncoding
+.. autoclass:: SinusoidalPositionalEncoding
+
+   
+
+   
+   .. rubric:: Methods
+
+   .. autosummary::
+   
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Softshrink.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Softshrink.rst
@@ -3,4 +3,14 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: Softshrink
+.. autoclass:: Softshrink
+
+   
+
+   
+   .. rubric:: Methods
+
+   .. autosummary::
+   
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Step.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Step.rst
@@ -3,4 +3,14 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: Step
+.. autoclass:: Step
+
+   
+
+   
+   .. rubric:: Methods
+
+   .. autosummary::
+   
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Transformer.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Transformer.rst
@@ -3,4 +3,14 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: Transformer
+.. autoclass:: Transformer
+
+   
+
+   
+   .. rubric:: Methods
+
+   .. autosummary::
+   
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Upsample.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Upsample.rst
@@ -3,4 +3,14 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: Upsample
+.. autoclass:: Upsample
+
+   
+
+   
+   .. rubric:: Methods
+
+   .. autosummary::
+   
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.elu.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.elu.rst
@@ -3,4 +3,9 @@

 .. currentmodule:: mlx.nn

-.. autofunction:: elu
+.. autoclass:: elu
+
+   
+
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.gelu.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.gelu.rst
@@ -3,4 +3,9 @@

 .. currentmodule:: mlx.nn

-.. autofunction:: gelu
+.. autoclass:: gelu
+
+   
+
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.gelu_approx.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.gelu_approx.rst
@@ -3,4 +3,9 @@

 .. currentmodule:: mlx.nn

-.. autofunction:: gelu_approx
+.. autoclass:: gelu_approx
+
+   
+
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.gelu_fast_approx.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.gelu_fast_approx.rst
@@ -3,4 +3,9 @@

 .. currentmodule:: mlx.nn

-.. autofunction:: gelu_fast_approx
+.. autoclass:: gelu_fast_approx
+
+   
+
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.glu.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.glu.rst
@@ -3,4 +3,9 @@

 .. currentmodule:: mlx.nn

-.. autofunction:: glu
+.. autoclass:: glu
+
+   
+
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.hardswish.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.hardswish.rst
@@ -3,4 +3,9 @@

 .. currentmodule:: mlx.nn

-.. autofunction:: hardswish
+.. autoclass:: hardswish
+
+   
+
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.leaky_relu.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.leaky_relu.rst
@@ -3,4 +3,9 @@

 .. currentmodule:: mlx.nn

-.. autofunction:: leaky_relu
+.. autoclass:: leaky_relu
+
+   
+
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.log_sigmoid.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.log_sigmoid.rst
@@ -3,4 +3,9 @@

 .. currentmodule:: mlx.nn

-.. autofunction:: log_sigmoid
+.. autoclass:: log_sigmoid
+
+   
+
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.log_softmax.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.log_softmax.rst
@@ -3,4 +3,9 @@

 .. currentmodule:: mlx.nn

-.. autofunction:: log_softmax
+.. autoclass:: log_softmax
+
+   
+
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.losses.binary_cross_entropy.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.losses.binary_cross_entropy.rst
@@ -3,4 +3,9 @@

 .. currentmodule:: mlx.nn.losses

-.. autofunction:: binary_cross_entropy
+.. autoclass:: binary_cross_entropy
+
+   
+
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.losses.cosine_similarity_loss.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.losses.cosine_similarity_loss.rst
@@ -3,4 +3,9 @@

 .. currentmodule:: mlx.nn.losses

-.. autofunction:: cosine_similarity_loss
+.. autoclass:: cosine_similarity_loss
+
+   
+
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.losses.cross_entropy.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.losses.cross_entropy.rst
@@ -3,4 +3,9 @@

 .. currentmodule:: mlx.nn.losses

-.. autofunction:: cross_entropy
+.. autoclass:: cross_entropy
+
+   
+
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.losses.gaussian_nll_loss.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.losses.gaussian_nll_loss.rst
@@ -3,4 +3,9 @@

 .. currentmodule:: mlx.nn.losses

-.. autofunction:: gaussian_nll_loss
+.. autoclass:: gaussian_nll_loss
+
+   
+
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.losses.hinge_loss.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.losses.hinge_loss.rst
@@ -3,4 +3,9 @@

 .. currentmodule:: mlx.nn.losses

-.. autofunction:: hinge_loss
+.. autoclass:: hinge_loss
+
+   
+
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.losses.huber_loss.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.losses.huber_loss.rst
@@ -3,4 +3,9 @@

 .. currentmodule:: mlx.nn.losses

-.. autofunction:: huber_loss
+.. autoclass:: huber_loss
+
+   
+
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.losses.kl_div_loss.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.losses.kl_div_loss.rst
@@ -3,4 +3,9 @@

 .. currentmodule:: mlx.nn.losses

-.. autofunction:: kl_div_loss
+.. autoclass:: kl_div_loss
+
+   
+
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.losses.l1_loss.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.losses.l1_loss.rst
@@ -3,4 +3,9 @@

 .. currentmodule:: mlx.nn.losses

-.. autofunction:: l1_loss
+.. autoclass:: l1_loss
+
+   
+
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.losses.log_cosh_loss.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.losses.log_cosh_loss.rst
@@ -3,4 +3,9 @@

 .. currentmodule:: mlx.nn.losses

-.. autofunction:: log_cosh_loss
+.. autoclass:: log_cosh_loss
+
+   
+
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.losses.margin_ranking_loss.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.losses.margin_ranking_loss.rst
@@ -3,4 +3,9 @@

 .. currentmodule:: mlx.nn.losses

-.. autofunction:: margin_ranking_loss
+.. autoclass:: margin_ranking_loss
+
+   
+
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.losses.mse_loss.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.losses.mse_loss.rst
@@ -3,4 +3,9 @@

 .. currentmodule:: mlx.nn.losses

-.. autofunction:: mse_loss
+.. autoclass:: mse_loss
+
+   
+
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.losses.nll_loss.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.losses.nll_loss.rst
@@ -3,4 +3,9 @@

 .. currentmodule:: mlx.nn.losses

-.. autofunction:: nll_loss
+.. autoclass:: nll_loss
+
+   
+
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.losses.smooth_l1_loss.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.losses.smooth_l1_loss.rst
@@ -3,4 +3,9 @@

 .. currentmodule:: mlx.nn.losses

-.. autofunction:: smooth_l1_loss
+.. autoclass:: smooth_l1_loss
+
+   
+
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.losses.triplet_loss.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.losses.triplet_loss.rst
@@ -3,4 +3,9 @@

 .. currentmodule:: mlx.nn.losses

-.. autofunction:: triplet_loss
+.. autoclass:: triplet_loss
+
+   
+
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.mish.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.mish.rst
@@ -3,4 +3,9 @@

 .. currentmodule:: mlx.nn

-.. autofunction:: mish
+.. autoclass:: mish
+
+   
+
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.prelu.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.prelu.rst
@@ -3,4 +3,9 @@

 .. currentmodule:: mlx.nn

-.. autofunction:: prelu
+.. autoclass:: prelu
+
+   
+
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.relu.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.relu.rst
@@ -3,4 +3,9 @@

 .. currentmodule:: mlx.nn

-.. autofunction:: relu
+.. autoclass:: relu
+
+   
+
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.relu6.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.relu6.rst
@@ -3,4 +3,9 @@

 .. currentmodule:: mlx.nn

-.. autofunction:: relu6
+.. autoclass:: relu6
+
+   
+
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.selu.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.selu.rst
@@ -3,4 +3,9 @@

 .. currentmodule:: mlx.nn

-.. autofunction:: selu
+.. autoclass:: selu
+
+   
+
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.sigmoid.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.sigmoid.rst
@@ -3,4 +3,9 @@

 .. currentmodule:: mlx.nn

-.. autofunction:: sigmoid
+.. autoclass:: sigmoid
+
+   
+
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.silu.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.silu.rst
@@ -3,4 +3,9 @@

 .. currentmodule:: mlx.nn

-.. autofunction:: silu
+.. autoclass:: silu
+
+   
+
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.softmax.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.softmax.rst
@@ -3,4 +3,9 @@

 .. currentmodule:: mlx.nn

-.. autofunction:: softmax
+.. autoclass:: softmax
+
+   
+
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.softplus.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.softplus.rst
@@ -3,4 +3,9 @@

 .. currentmodule:: mlx.nn

-.. autofunction:: softplus
+.. autoclass:: softplus
+
+   
+
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.softshrink.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.softshrink.rst
@@ -3,4 +3,9 @@

 .. currentmodule:: mlx.nn

-.. autofunction:: softshrink
+.. autoclass:: softshrink
+
+   
+
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.step.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.step.rst
@@ -3,4 +3,9 @@

 .. currentmodule:: mlx.nn

-.. autofunction:: step
+.. autoclass:: step
+
+   
+
+   
+   
--- a/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.tanh.rst
+++ b/docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.tanh.rst
@@ -3,4 +3,9 @@

 .. currentmodule:: mlx.nn

-.. autofunction:: tanh
+.. autoclass:: tanh
+
+   
+
+   
+   
--- a/docs/build/html/_sources/python/nn/layers.rst
+++ b/docs/build/html/_sources/python/nn/layers.rst
@@ -31,6 +31,7 @@ Layers
   Mish
   MultiHeadAttention
   PReLU
+   QuantizedEmbedding
   QuantizedLinear
   RMSNorm
   ReLU
@@ -43,4 +44,4 @@ Layers
   Softshrink
   Step
   Transformer
-   Upsample
+   Upsample