docs update

2025-10-17 19:14:55 +00:00
parent b8691a1637
commit 3ee92b5d61
437 changed files with 11568 additions and 13689 deletions
@@ -0,0 +1,6 @@
+mlx.core.compile
+================
+
+.. currentmodule:: mlx.core
+
+.. autofunction:: compile
@@ -0,0 +1,6 @@
+mlx.core.disable\_compile
+=========================
+
+.. currentmodule:: mlx.core
+
+.. autofunction:: disable_compile
@@ -0,0 +1,6 @@
+mlx.core.enable\_compile
+========================
+
+.. currentmodule:: mlx.core
+
+.. autofunction:: enable_compile
@@ -1,6 +0,0 @@
-mlx.core.simplify
-=================
-
-.. currentmodule:: mlx.core
-
-.. autofunction:: simplify
@@ -14,5 +14,6 @@
   
         ~AdaDelta.__init__
         ~AdaDelta.apply_single
+         ~AdaDelta.init_single
   
   
@@ -14,5 +14,6 @@
   
         ~Adafactor.__init__
         ~Adafactor.apply_single
+         ~Adafactor.init_single
   
   
@@ -14,5 +14,6 @@
   
         ~Adagrad.__init__
         ~Adagrad.apply_single
+         ~Adagrad.init_single
   
   
@@ -14,5 +14,6 @@
   
         ~Adam.__init__
         ~Adam.apply_single
+         ~Adam.init_single
   
   
@@ -14,5 +14,6 @@
   
         ~Adamax.__init__
         ~Adamax.apply_single
+         ~Adamax.init_single
   
   
@@ -14,5 +14,6 @@
   
         ~Lion.__init__
         ~Lion.apply_single
+         ~Lion.init_single
   
   
@@ -0,0 +1,6 @@
+mlx.optimizers.Optimizer.apply\_gradients
+=========================================
+
+.. currentmodule:: mlx.optimizers
+
+.. automethod:: Optimizer.apply_gradients
@@ -0,0 +1,6 @@
+mlx.optimizers.Optimizer.init
+=============================
+
+.. currentmodule:: mlx.optimizers
+
+.. automethod:: Optimizer.init
@@ -1,20 +0,0 @@
-mlx.optimizers.Optimizer
-========================
-
-.. currentmodule:: mlx.optimizers
-
-.. autoclass:: Optimizer
-
-   
-
-   
-   .. rubric:: Methods
-
-   .. autosummary::
-   
-         ~Optimizer.__init__
-         ~Optimizer.apply_gradients
-         ~Optimizer.apply_single
-         ~Optimizer.update
-   
-   
@@ -0,0 +1,6 @@
+mlx.optimizers.Optimizer.state
+==============================
+
+.. currentmodule:: mlx.optimizers
+
+.. autoproperty:: Optimizer.state
@@ -0,0 +1,6 @@
+mlx.optimizers.Optimizer.update
+===============================
+
+.. currentmodule:: mlx.optimizers
+
+.. automethod:: Optimizer.update
@@ -1,17 +0,0 @@
-mlx.optimizers.OptimizerState
-=============================
-
-.. currentmodule:: mlx.optimizers
-
-.. autoclass:: OptimizerState
-
-   
-
-   
-   .. rubric:: Methods
-
-   .. autosummary::
-   
-         ~OptimizerState.get
-   
-   
@@ -14,5 +14,6 @@
   
         ~RMSprop.__init__
         ~RMSprop.apply_single
+         ~RMSprop.init_single
   
   
@@ -14,5 +14,6 @@
   
         ~SGD.__init__
         ~SGD.apply_single
+         ~SGD.init_single
   
   
@@ -3,6 +3,4 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: ALiBi
-
-   
+.. autoclass:: ALiBi
@@ -3,6 +3,4 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: BatchNorm
-
-   
+.. autoclass:: BatchNorm
@@ -3,6 +3,4 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: Conv1d
-
-   
+.. autoclass:: Conv1d
@@ -3,6 +3,4 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: Conv2d
-
-   
+.. autoclass:: Conv2d
@@ -3,6 +3,4 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: Dropout
-
-   
+.. autoclass:: Dropout
@@ -3,6 +3,4 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: Dropout2d
-
-   
+.. autoclass:: Dropout2d
@@ -3,6 +3,4 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: Dropout3d
-
-   
+.. autoclass:: Dropout3d
@@ -3,6 +3,4 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: Embedding
-
-   
+.. autoclass:: Embedding
@@ -3,6 +3,4 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: GELU
-
-   
+.. autoclass:: GELU
@@ -3,6 +3,4 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: GroupNorm
-
-   
+.. autoclass:: GroupNorm
@@ -3,6 +3,4 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: InstanceNorm
-
-   
+.. autoclass:: InstanceNorm
@@ -3,6 +3,4 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: LayerNorm
-
-   
+.. autoclass:: LayerNorm
@@ -3,6 +3,4 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: Linear
-
-   
+.. autoclass:: Linear
@@ -3,6 +3,4 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: Mish
-
-   
+.. autoclass:: Mish
@@ -0,0 +1,6 @@
+mlx.nn.Module.state
+===================
+
+.. currentmodule:: mlx.nn
+
+.. autoproperty:: Module.state
@@ -3,6 +3,4 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: MultiHeadAttention
-
-   
+.. autoclass:: MultiHeadAttention
@@ -3,6 +3,4 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: PReLU
-
-   
+.. autoclass:: PReLU
@@ -3,6 +3,4 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: QuantizedLinear
-
-   
+.. autoclass:: QuantizedLinear
@@ -3,6 +3,4 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: RMSNorm
-
-   
+.. autoclass:: RMSNorm
@@ -3,6 +3,4 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: ReLU
-
-   
+.. autoclass:: ReLU
@@ -3,6 +3,4 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: RoPE
-
-   
+.. autoclass:: RoPE
@@ -3,6 +3,4 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: SELU
-
-   
+.. autoclass:: SELU
@@ -3,6 +3,4 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: Sequential
-
-   
+.. autoclass:: Sequential
@@ -3,6 +3,4 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: SiLU
-
-   
+.. autoclass:: SiLU
@@ -3,6 +3,4 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: SinusoidalPositionalEncoding
-
-   
+.. autoclass:: SinusoidalPositionalEncoding
@@ -3,6 +3,4 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: Softshrink
-
-   
+.. autoclass:: Softshrink
@@ -3,6 +3,4 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: Step
-
-   
+.. autoclass:: Step
@@ -3,6 +3,4 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: Transformer
-
-   
+.. autoclass:: Transformer
@@ -3,6 +3,4 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: gelu
-
-   
+.. autofunction:: gelu
@@ -3,6 +3,4 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: gelu_approx
-
-   
+.. autofunction:: gelu_approx
@@ -3,6 +3,4 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: gelu_fast_approx
-
-   
+.. autofunction:: gelu_fast_approx
@@ -1,6 +0,0 @@
-mlx.nn.init.constant
-====================
-
-.. currentmodule:: mlx.nn.init
-
-.. autofunction:: constant
@@ -1,6 +0,0 @@
-mlx.nn.init.glorot\_normal
-==========================
-
-.. currentmodule:: mlx.nn.init
-
-.. autofunction:: glorot_normal
@@ -1,6 +0,0 @@
-mlx.nn.init.glorot\_uniform
-===========================
-
-.. currentmodule:: mlx.nn.init
-
-.. autofunction:: glorot_uniform
@@ -1,6 +0,0 @@
-mlx.nn.init.he\_normal
-======================
-
-.. currentmodule:: mlx.nn.init
-
-.. autofunction:: he_normal
@@ -1,6 +0,0 @@
-mlx.nn.init.he\_uniform
-=======================
-
-.. currentmodule:: mlx.nn.init
-
-.. autofunction:: he_uniform
@@ -1,6 +0,0 @@
-mlx.nn.init.identity
-====================
-
-.. currentmodule:: mlx.nn.init
-
-.. autofunction:: identity
@@ -1,6 +0,0 @@
-mlx.nn.init.normal
-==================
-
-.. currentmodule:: mlx.nn.init
-
-.. autofunction:: normal
@@ -1,6 +0,0 @@
-mlx.nn.init.uniform
-===================
-
-.. currentmodule:: mlx.nn.init
-
-.. autofunction:: uniform
@@ -1,6 +0,0 @@
-mlx.nn.initializers.constant
-============================
-
-.. currentmodule:: mlx.nn.initializers
-
-.. autofunction:: constant
@@ -1,6 +0,0 @@
-mlx.nn.initializers.glorot\_normal
-==================================
-
-.. currentmodule:: mlx.nn.initializers
-
-.. autofunction:: glorot_normal
@@ -1,6 +0,0 @@
-mlx.nn.initializers.glorot\_uniform
-===================================
-
-.. currentmodule:: mlx.nn.initializers
-
-.. autofunction:: glorot_uniform
@@ -1,6 +0,0 @@
-mlx.nn.initializers.he\_normal
-==============================
-
-.. currentmodule:: mlx.nn.initializers
-
-.. autofunction:: he_normal
@@ -1,6 +0,0 @@
-mlx.nn.initializers.he\_uniform
-===============================
-
-.. currentmodule:: mlx.nn.initializers
-
-.. autofunction:: he_uniform
@@ -1,6 +0,0 @@
-mlx.nn.initializers.identity
-============================
-
-.. currentmodule:: mlx.nn.initializers
-
-.. autofunction:: identity
@@ -1,6 +0,0 @@
-mlx.nn.initializers.normal
-==========================
-
-.. currentmodule:: mlx.nn.initializers
-
-.. autofunction:: normal
@@ -1,6 +0,0 @@
-mlx.nn.initializers.uniform
-===========================
-
-.. currentmodule:: mlx.nn.initializers
-
-.. autofunction:: uniform
@@ -3,6 +3,4 @@

 .. currentmodule:: mlx.nn.losses

-.. autoclass:: binary_cross_entropy
-
-   
+.. autofunction:: binary_cross_entropy
@@ -3,6 +3,4 @@

 .. currentmodule:: mlx.nn.losses

-.. autoclass:: cosine_similarity_loss
-
-   
+.. autofunction:: cosine_similarity_loss
@@ -3,6 +3,4 @@

 .. currentmodule:: mlx.nn.losses

-.. autoclass:: cross_entropy
-
-   
+.. autofunction:: cross_entropy
@@ -3,6 +3,4 @@

 .. currentmodule:: mlx.nn.losses

-.. autoclass:: gaussian_nll_loss
-
-   
+.. autofunction:: gaussian_nll_loss
@@ -3,6 +3,4 @@

 .. currentmodule:: mlx.nn.losses

-.. autoclass:: hinge_loss
-
-   
+.. autofunction:: hinge_loss
@@ -3,6 +3,4 @@

 .. currentmodule:: mlx.nn.losses

-.. autoclass:: huber_loss
-
-   
+.. autofunction:: huber_loss
@@ -3,6 +3,4 @@

 .. currentmodule:: mlx.nn.losses

-.. autoclass:: kl_div_loss
-
-   
+.. autofunction:: kl_div_loss
@@ -3,6 +3,4 @@

 .. currentmodule:: mlx.nn.losses

-.. autoclass:: l1_loss
-
-   
+.. autofunction:: l1_loss
@@ -3,6 +3,4 @@

 .. currentmodule:: mlx.nn.losses

-.. autoclass:: log_cosh_loss
-
-   
+.. autofunction:: log_cosh_loss
@@ -0,0 +1,6 @@
+mlx.nn.losses.margin\_ranking\_loss
+===================================
+
+.. currentmodule:: mlx.nn.losses
+
+.. autofunction:: margin_ranking_loss
@@ -3,6 +3,4 @@

 .. currentmodule:: mlx.nn.losses

-.. autoclass:: mse_loss
-
-   
+.. autofunction:: mse_loss
@@ -3,6 +3,4 @@

 .. currentmodule:: mlx.nn.losses

-.. autoclass:: nll_loss
-
-   
+.. autofunction:: nll_loss
@@ -3,6 +3,4 @@

 .. currentmodule:: mlx.nn.losses

-.. autoclass:: smooth_l1_loss
-
-   
+.. autofunction:: smooth_l1_loss
@@ -3,6 +3,4 @@

 .. currentmodule:: mlx.nn.losses

-.. autoclass:: triplet_loss
-
-   
+.. autofunction:: triplet_loss
@@ -3,6 +3,4 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: mish
-
-   
+.. autofunction:: mish
@@ -3,6 +3,4 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: prelu
-
-   
+.. autofunction:: prelu
@@ -3,6 +3,4 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: relu
-
-   
+.. autofunction:: relu
@@ -3,6 +3,4 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: selu
-
-   
+.. autofunction:: selu
@@ -3,6 +3,4 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: silu
-
-   
+.. autofunction:: silu
@@ -3,6 +3,4 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: softshrink
-
-   
+.. autofunction:: softshrink
@@ -3,6 +3,4 @@

 .. currentmodule:: mlx.nn

-.. autoclass:: step
-
-   
+.. autofunction:: step
@@ -1,18 +0,0 @@
-.. _initializers:
-
-.. currentmodule:: mlx.nn.initializers
-
-Initializers
--------------
-
-.. autosummary::
-   :toctree: _autosummary_functions
-
-   constant
-   normal
-   uniform
-   identity
-   glorot_normal
-   glorot_uniform
-   he_normal
-   he_uniform
@@ -18,6 +18,7 @@ Loss Functions
   kl_div_loss
   l1_loss
   log_cosh_loss
+   margin_ranking_loss
   mse_loss
   nll_loss
   smooth_l1_loss
@@ -11,6 +11,7 @@ Module
      :toctree: _autosummary
   
      Module.training
+      Module.state
   
   .. rubric:: Methods

@@ -0,0 +1,23 @@
+Optimizer
+=========
+
+.. currentmodule:: mlx.optimizers
+
+.. autoclass:: Optimizer 
+
+
+   .. rubric:: Attributes
+
+   .. autosummary::
+      :toctree: _autosummary
+
+      Optimizer.state
+   
+   .. rubric:: Methods
+
+   .. autosummary::
+      :toctree: _autosummary
+   
+      Optimizer.apply_gradients
+      Optimizer.init
+      Optimizer.update
@@ -29,14 +29,16 @@ model's parameters and the **optimizer state**.
            # Compute the new parameters but also the optimizer state.
            mx.eval(model.parameters(), optimizer.state)

+.. toctree::
+
+   optimizer
+
 .. currentmodule:: mlx.optimizers

 .. autosummary::
   :toctree: _autosummary
   :template: optimizers-template.rst

-   OptimizerState
-   Optimizer
   SGD
   RMSprop
   Adagrad
@@ -9,6 +9,9 @@ Transforms
  :toctree: _autosummary

   eval
+   compile
+   disable_compile
+   enable_compile
   grad
   value_and_grad
   jvp