docs update

2025-09-22 21:48:08 +08:00 · 2024-04-26 08:24:09 -07:00
parent f77d99b285
commit f946f689a6
502 changed files with 14665 additions and 15103 deletions
--- a/docs/build/html/_sources/python/_autosummary/mlx.core.block_masked_mm.rst
+++ b/docs/build/html/_sources/python/_autosummary/mlx.core.block_masked_mm.rst
@@ -0,0 +1,6 @@
+mlx.core.block\_masked\_mm
+==========================
+
+.. currentmodule:: mlx.core
+
+.. autofunction:: block_masked_mm
--- a/docs/build/html/_sources/python/_autosummary/mlx.core.degrees.rst
+++ b/docs/build/html/_sources/python/_autosummary/mlx.core.degrees.rst
@@ -0,0 +1,6 @@
+mlx.core.degrees
+================
+
+.. currentmodule:: mlx.core
+
+.. autofunction:: degrees
--- a/docs/build/html/_sources/python/_autosummary/mlx.core.metal.clear_cache.rst
+++ b/docs/build/html/_sources/python/_autosummary/mlx.core.metal.clear_cache.rst
@@ -0,0 +1,6 @@
+mlx.core.metal.clear\_cache
+===========================
+
+.. currentmodule:: mlx.core.metal
+
+.. autofunction:: clear_cache
--- a/docs/build/html/_sources/python/_autosummary/mlx.core.not_equal.rst
+++ b/docs/build/html/_sources/python/_autosummary/mlx.core.not_equal.rst
@@ -0,0 +1,6 @@
+mlx.core.not\_equal
+===================
+
+.. currentmodule:: mlx.core
+
+.. autofunction:: not_equal
--- a/docs/build/html/_sources/python/_autosummary/mlx.core.radians.rst
+++ b/docs/build/html/_sources/python/_autosummary/mlx.core.radians.rst
@@ -0,0 +1,6 @@
+mlx.core.radians
+================
+
+.. currentmodule:: mlx.core
+
+.. autofunction:: radians
--- a/docs/build/html/_sources/python/_autosummary/mlx.core.synchronize.rst
+++ b/docs/build/html/_sources/python/_autosummary/mlx.core.synchronize.rst
@@ -0,0 +1,6 @@
+mlx.core.synchronize
+====================
+
+.. currentmodule:: mlx.core
+
+.. autofunction:: synchronize
--- a/docs/build/html/_sources/python/_autosummary/mlx.nn.quantize.rst
+++ b/docs/build/html/_sources/python/_autosummary/mlx.nn.quantize.rst
@@ -0,0 +1,6 @@
+mlx.nn.quantize
+===============
+
+.. currentmodule:: mlx.nn
+
+.. autofunction:: quantize
--- a/docs/build/html/_sources/python/_autosummary/mlx.utils.tree_map_with_path.rst
+++ b/docs/build/html/_sources/python/_autosummary/mlx.utils.tree_map_with_path.rst
@@ -0,0 +1,6 @@
+mlx.utils.tree\_map\_with\_path
+===============================
+
+.. currentmodule:: mlx.utils
+
+.. autofunction:: tree_map_with_path
--- a/docs/build/html/_sources/python/devices_and_streams.rst
+++ b/docs/build/html/_sources/python/devices_and_streams.rst
@@ -16,3 +16,4 @@ Devices and Streams
   new_stream
   set_default_stream
   stream
+   synchronize
--- a/docs/build/html/_sources/python/metal.rst
+++ b/docs/build/html/_sources/python/metal.rst
@@ -12,5 +12,6 @@ Metal
  get_cache_memory
  set_memory_limit
  set_cache_limit
+  clear_cache
  start_capture
  stop_capture
--- a/docs/build/html/_sources/python/nn.rst
+++ b/docs/build/html/_sources/python/nn.rst
@@ -173,6 +173,7 @@ In detail:
   :toctree: _autosummary

   value_and_grad
+   quantize

 .. toctree::

--- 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
--- a/docs/build/html/_sources/python/ops.rst
+++ b/docs/build/html/_sources/python/ops.rst
@@ -29,6 +29,7 @@ Operations
   atleast_2d
   atleast_3d
   broadcast_to
+   block_masked_mm
   ceil
   clip
   concatenate
@@ -42,6 +43,7 @@ Operations
   cummin
   cumprod
   cumsum
+   degrees
   dequantize
   diag
   diagonal
@@ -90,6 +92,7 @@ Operations
   moveaxis
   multiply
   negative
+   not_equal
   ones
   ones_like
   outer
@@ -98,6 +101,7 @@ Operations
   prod
   quantize
   quantized_matmul
+   radians
   reciprocal
   repeat
   reshape
--- a/docs/build/html/_sources/python/tree_utils.rst
+++ b/docs/build/html/_sources/python/tree_utils.rst
@@ -19,3 +19,4 @@ return python trees will be using the default python ``dict``, ``list`` and
   tree_flatten
   tree_unflatten
   tree_map
+   tree_map_with_path