docs update

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