docs update

docs/build/html/_sources/dev/extensions.rst

@@ -35,7 +35,7 @@ However, you work with vector math libraries often and realize that the
 You would really like the part of your applications that does this operation 
 on the CPU to be very fast - so you decide that you want it to rely on the 
 ``axpby`` routine provided by the Accelerate_ framework. Continuing to impose 
-our assumptions on to you, let's also assume that you want to learn how add 
+our assumptions on to you, let's also assume that you want to learn how to add 
 your own implementation for the gradients of your new operation while going 
 over the ins-and-outs of the MLX framework. 
 

docs/build/html/_sources/python/_autosummary/mlx.core.Stream.rst

@@ -1,28 +1,6 @@
-mlx.core.Stream
+mlx.core.stream
 ===============
 
 .. currentmodule:: mlx.core
 
-.. autoclass:: Stream
-
-   
-   .. automethod:: __init__
-
-   
-   .. rubric:: Methods
-
-   .. autosummary::
-   
-      ~Stream.__init__
-   
-   
-
-   
-   
-   .. rubric:: Attributes
-
-   .. autosummary::
-   
-      ~Stream.device
-   
-   
+.. autofunction:: stream

docs/build/html/_sources/python/_autosummary/stream_class.rst

@@ -0,0 +1,28 @@
+mlx.core.Stream
+===============
+
+.. currentmodule:: mlx.core
+
+.. autoclass:: Stream
+
+   
+   .. automethod:: __init__
+
+   
+   .. rubric:: Methods
+
+   .. autosummary::
+   
+      ~Stream.__init__
+   
+   
+
+   
+   
+   .. rubric:: Attributes
+
+   .. autosummary::
+   
+      ~Stream.device
+   
+   

docs/build/html/_sources/python/_autosummary_stream/Stream.rst

@@ -0,0 +1,28 @@
+mlx.core.Stream
+===============
+
+.. currentmodule:: mlx.core
+
+.. autoclass:: Stream
+
+   
+   .. automethod:: __init__
+
+   
+   .. rubric:: Methods
+
+   .. autosummary::
+   
+      ~Stream.__init__
+   
+   
+
+   
+   
+   .. rubric:: Attributes
+
+   .. autosummary::
+   
+      ~Stream.device
+   
+   

docs/build/html/_sources/python/_autosummary_stream/mlx.core.Device.rst

@@ -0,0 +1,28 @@
+mlx.core.Device
+===============
+
+.. currentmodule:: mlx.core
+
+.. autoclass:: Device
+
+   
+   .. automethod:: __init__
+
+   
+   .. rubric:: Methods
+
+   .. autosummary::
+   
+      ~Device.__init__
+   
+   
+
+   
+   
+   .. rubric:: Attributes
+
+   .. autosummary::
+   
+      ~Device.type
+   
+   

docs/build/html/_sources/python/_autosummary_stream/mlx.core.Stream.rst

@@ -0,0 +1,28 @@
+mlx.core.Stream
+===============
+
+.. currentmodule:: mlx.core
+
+.. autoclass:: Stream
+
+   
+   .. automethod:: __init__
+
+   
+   .. rubric:: Methods
+
+   .. autosummary::
+   
+      ~Stream.__init__
+   
+   
+
+   
+   
+   .. rubric:: Attributes
+
+   .. autosummary::
+   
+      ~Stream.device
+   
+   

docs/build/html/_sources/python/devices_and_streams.rst

@@ -9,9 +9,10 @@ Devices and Streams
   :toctree: _autosummary
 
    Device
+   Stream
    default_device
    set_default_device
-   Stream
    default_stream
    new_stream
    set_default_stream
+   stream

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

@@ -0,0 +1,6 @@
+mlx.nn.AvgPool1d
+================
+
+.. currentmodule:: mlx.nn
+
+.. autoclass:: AvgPool1d

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

@@ -0,0 +1,6 @@
+mlx.nn.AvgPool2d
+================
+
+.. currentmodule:: mlx.nn
+
+.. autoclass:: AvgPool2d

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

@@ -0,0 +1,6 @@
+mlx.nn.MaxPool1d
+================
+
+.. currentmodule:: mlx.nn
+
+.. autoclass:: MaxPool1d

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

@@ -0,0 +1,6 @@
+mlx.nn.MaxPool2d
+================
+
+.. currentmodule:: mlx.nn
+
+.. autoclass:: MaxPool2d

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

@@ -10,6 +10,8 @@ Layers
    :template: nn-module-template.rst
 
    ALiBi
+   AvgPool1d
+   AvgPool2d
    BatchNorm
    Conv1d
    Conv2d
@@ -22,6 +24,8 @@ Layers
    InstanceNorm
    LayerNorm
    Linear
+   MaxPool1d
+   MaxPool2d
    Mish
    MultiHeadAttention
    PReLU

docs/build/html/_sources/python/optimizers.rst

@@ -31,20 +31,6 @@ model's parameters and the **optimizer state**.
 
 .. toctree::
 
-   optimizer
-
-.. currentmodule:: mlx.optimizers
-
-.. autosummary::
-   :toctree: _autosummary
-   :template: optimizers-template.rst
-
-   SGD
-   RMSprop
-   Adagrad
-   Adafactor
-   AdaDelta
-   Adam
-   AdamW
-   Adamax
-   Lion
+   optimizers/optimizer
+   optimizers/common_optimizers
+   optimizers/schedulers

docs/build/html/_sources/python/optimizers/_autosummary/mlx.optimizers.AdaDelta.rst

@@ -0,0 +1,19 @@
+mlx.optimizers.AdaDelta
+=======================
+
+.. currentmodule:: mlx.optimizers
+
+.. autoclass:: AdaDelta
+
+   
+
+   
+   .. rubric:: Methods
+
+   .. autosummary::
+   
+         ~AdaDelta.__init__
+         ~AdaDelta.apply_single
+         ~AdaDelta.init_single
+   
+   

docs/build/html/_sources/python/optimizers/_autosummary/mlx.optimizers.Adafactor.rst

@@ -0,0 +1,19 @@
+mlx.optimizers.Adafactor
+========================
+
+.. currentmodule:: mlx.optimizers
+
+.. autoclass:: Adafactor
+
+   
+
+   
+   .. rubric:: Methods
+
+   .. autosummary::
+   
+         ~Adafactor.__init__
+         ~Adafactor.apply_single
+         ~Adafactor.init_single
+   
+   

docs/build/html/_sources/python/optimizers/_autosummary/mlx.optimizers.Adagrad.rst

@@ -0,0 +1,19 @@
+mlx.optimizers.Adagrad
+======================
+
+.. currentmodule:: mlx.optimizers
+
+.. autoclass:: Adagrad
+
+   
+
+   
+   .. rubric:: Methods
+
+   .. autosummary::
+   
+         ~Adagrad.__init__
+         ~Adagrad.apply_single
+         ~Adagrad.init_single
+   
+   

docs/build/html/_sources/python/optimizers/_autosummary/mlx.optimizers.Adam.rst

@@ -0,0 +1,19 @@
+mlx.optimizers.Adam
+===================
+
+.. currentmodule:: mlx.optimizers
+
+.. autoclass:: Adam
+
+   
+
+   
+   .. rubric:: Methods
+
+   .. autosummary::
+   
+         ~Adam.__init__
+         ~Adam.apply_single
+         ~Adam.init_single
+   
+   

docs/build/html/_sources/python/optimizers/_autosummary/mlx.optimizers.AdamW.rst

@@ -0,0 +1,18 @@
+mlx.optimizers.AdamW
+====================
+
+.. currentmodule:: mlx.optimizers
+
+.. autoclass:: AdamW
+
+   
+
+   
+   .. rubric:: Methods
+
+   .. autosummary::
+   
+         ~AdamW.__init__
+         ~AdamW.apply_single
+   
+   

docs/build/html/_sources/python/optimizers/_autosummary/mlx.optimizers.Adamax.rst

@@ -0,0 +1,19 @@
+mlx.optimizers.Adamax
+=====================
+
+.. currentmodule:: mlx.optimizers
+
+.. autoclass:: Adamax
+
+   
+
+   
+   .. rubric:: Methods
+
+   .. autosummary::
+   
+         ~Adamax.__init__
+         ~Adamax.apply_single
+         ~Adamax.init_single
+   
+   

docs/build/html/_sources/python/optimizers/_autosummary/mlx.optimizers.Lion.rst

@@ -0,0 +1,19 @@
+mlx.optimizers.Lion
+===================
+
+.. currentmodule:: mlx.optimizers
+
+.. autoclass:: Lion
+
+   
+
+   
+   .. rubric:: Methods
+
+   .. autosummary::
+   
+         ~Lion.__init__
+         ~Lion.apply_single
+         ~Lion.init_single
+   
+   

docs/build/html/_sources/python/optimizers/_autosummary/mlx.optimizers.Optimizer.apply_gradients.rst

@@ -0,0 +1,6 @@
+mlx.optimizers.Optimizer.apply\_gradients
+=========================================
+
+.. currentmodule:: mlx.optimizers
+
+.. automethod:: Optimizer.apply_gradients

docs/build/html/_sources/python/optimizers/_autosummary/mlx.optimizers.Optimizer.init.rst

@@ -0,0 +1,6 @@
+mlx.optimizers.Optimizer.init
+=============================
+
+.. currentmodule:: mlx.optimizers
+
+.. automethod:: Optimizer.init

docs/build/html/_sources/python/optimizers/_autosummary/mlx.optimizers.Optimizer.state.rst

@@ -0,0 +1,6 @@
+mlx.optimizers.Optimizer.state
+==============================
+
+.. currentmodule:: mlx.optimizers
+
+.. autoproperty:: Optimizer.state

docs/build/html/_sources/python/optimizers/_autosummary/mlx.optimizers.Optimizer.update.rst

@@ -0,0 +1,6 @@
+mlx.optimizers.Optimizer.update
+===============================
+
+.. currentmodule:: mlx.optimizers
+
+.. automethod:: Optimizer.update

docs/build/html/_sources/python/optimizers/_autosummary/mlx.optimizers.RMSprop.rst

@@ -0,0 +1,19 @@
+mlx.optimizers.RMSprop
+======================
+
+.. currentmodule:: mlx.optimizers
+
+.. autoclass:: RMSprop
+
+   
+
+   
+   .. rubric:: Methods
+
+   .. autosummary::
+   
+         ~RMSprop.__init__
+         ~RMSprop.apply_single
+         ~RMSprop.init_single
+   
+   

docs/build/html/_sources/python/optimizers/_autosummary/mlx.optimizers.SGD.rst

@@ -0,0 +1,19 @@
+mlx.optimizers.SGD
+==================
+
+.. currentmodule:: mlx.optimizers
+
+.. autoclass:: SGD
+
+   
+
+   
+   .. rubric:: Methods
+
+   .. autosummary::
+   
+         ~SGD.__init__
+         ~SGD.apply_single
+         ~SGD.init_single
+   
+   

docs/build/html/_sources/python/optimizers/_autosummary/mlx.optimizers.cosine_decay.rst

@@ -0,0 +1,6 @@
+mlx.optimizers.cosine\_decay
+============================
+
+.. currentmodule:: mlx.optimizers
+
+.. autofunction:: cosine_decay

docs/build/html/_sources/python/optimizers/_autosummary/mlx.optimizers.exponential_decay.rst

@@ -0,0 +1,6 @@
+mlx.optimizers.exponential\_decay
+=================================
+
+.. currentmodule:: mlx.optimizers
+
+.. autofunction:: exponential_decay

docs/build/html/_sources/python/optimizers/_autosummary/mlx.optimizers.step_decay.rst

@@ -0,0 +1,6 @@
+mlx.optimizers.step\_decay
+==========================
+
+.. currentmodule:: mlx.optimizers
+
+.. autofunction:: step_decay

docs/build/html/_sources/python/optimizers/common_optimizers.rst

@@ -0,0 +1,20 @@
+.. _common_optimizers:
+
+Common Optimizers
+=================
+
+.. currentmodule:: mlx.optimizers
+
+.. autosummary::
+   :toctree: _autosummary
+   :template: optimizers-template.rst
+
+   SGD
+   RMSprop
+   Adagrad
+   Adafactor
+   AdaDelta
+   Adam
+   AdamW
+   Adamax
+   Lion

docs/build/html/_sources/python/optimizers/optimizer.rst

@@ -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

docs/build/html/_sources/python/optimizers/schedulers.rst

@@ -0,0 +1,13 @@
+.. _schedulers:
+
+Schedulers
+==========
+
+.. currentmodule:: mlx.optimizers
+
+.. autosummary::
+   :toctree: _autosummary
+
+   step_decay    
+   exponential_decay    
+   cosine_decay