update docs

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

@@ -0,0 +1,8 @@
+mlx.nn.Softshrink
+=================
+
+.. currentmodule:: mlx.nn
+
+.. autoclass:: Softshrink
+
+   

docs/build/html/_sources/python/nn/_autosummary/mlx.nn.init.constant.rst

@@ -0,0 +1,6 @@
+mlx.nn.init.constant
+====================
+
+.. currentmodule:: mlx.nn.init
+
+.. autofunction:: constant

docs/build/html/_sources/python/nn/_autosummary/mlx.nn.init.glorot_normal.rst

@@ -0,0 +1,6 @@
+mlx.nn.init.glorot\_normal
+==========================
+
+.. currentmodule:: mlx.nn.init
+
+.. autofunction:: glorot_normal

docs/build/html/_sources/python/nn/_autosummary/mlx.nn.init.glorot_uniform.rst

@@ -0,0 +1,6 @@
+mlx.nn.init.glorot\_uniform
+===========================
+
+.. currentmodule:: mlx.nn.init
+
+.. autofunction:: glorot_uniform

docs/build/html/_sources/python/nn/_autosummary/mlx.nn.init.he_normal.rst

@@ -0,0 +1,6 @@
+mlx.nn.init.he\_normal
+======================
+
+.. currentmodule:: mlx.nn.init
+
+.. autofunction:: he_normal

docs/build/html/_sources/python/nn/_autosummary/mlx.nn.init.he_uniform.rst

@@ -0,0 +1,6 @@
+mlx.nn.init.he\_uniform
+=======================
+
+.. currentmodule:: mlx.nn.init
+
+.. autofunction:: he_uniform

docs/build/html/_sources/python/nn/_autosummary/mlx.nn.init.identity.rst

@@ -0,0 +1,6 @@
+mlx.nn.init.identity
+====================
+
+.. currentmodule:: mlx.nn.init
+
+.. autofunction:: identity

docs/build/html/_sources/python/nn/_autosummary/mlx.nn.init.normal.rst

@@ -0,0 +1,6 @@
+mlx.nn.init.normal
+==================
+
+.. currentmodule:: mlx.nn.init
+
+.. autofunction:: normal

docs/build/html/_sources/python/nn/_autosummary/mlx.nn.init.uniform.rst

@@ -0,0 +1,6 @@
+mlx.nn.init.uniform
+===================
+
+.. currentmodule:: mlx.nn.init
+
+.. autofunction:: uniform

docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.init.constant.rst

@@ -0,0 +1,6 @@
+mlx.nn.init.constant
+====================
+
+.. currentmodule:: mlx.nn.init
+
+.. autofunction:: constant

docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.init.glorot_normal.rst

@@ -0,0 +1,6 @@
+mlx.nn.init.glorot\_normal
+==========================
+
+.. currentmodule:: mlx.nn.init
+
+.. autofunction:: glorot_normal

docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.init.glorot_uniform.rst

@@ -0,0 +1,6 @@
+mlx.nn.init.glorot\_uniform
+===========================
+
+.. currentmodule:: mlx.nn.init
+
+.. autofunction:: glorot_uniform

docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.init.he_normal.rst

@@ -0,0 +1,6 @@
+mlx.nn.init.he\_normal
+======================
+
+.. currentmodule:: mlx.nn.init
+
+.. autofunction:: he_normal

docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.init.he_uniform.rst

@@ -0,0 +1,6 @@
+mlx.nn.init.he\_uniform
+=======================
+
+.. currentmodule:: mlx.nn.init
+
+.. autofunction:: he_uniform

docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.init.identity.rst

@@ -0,0 +1,6 @@
+mlx.nn.init.identity
+====================
+
+.. currentmodule:: mlx.nn.init
+
+.. autofunction:: identity

docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.init.normal.rst

@@ -0,0 +1,6 @@
+mlx.nn.init.normal
+==================
+
+.. currentmodule:: mlx.nn.init
+
+.. autofunction:: normal

docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.init.uniform.rst

@@ -0,0 +1,6 @@
+mlx.nn.init.uniform
+===================
+
+.. currentmodule:: mlx.nn.init
+
+.. autofunction:: uniform

docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.initializers.constant.rst

@@ -0,0 +1,6 @@
+mlx.nn.initializers.constant
+============================
+
+.. currentmodule:: mlx.nn.initializers
+
+.. autofunction:: constant

docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.initializers.glorot_normal.rst

@@ -0,0 +1,6 @@
+mlx.nn.initializers.glorot\_normal
+==================================
+
+.. currentmodule:: mlx.nn.initializers
+
+.. autofunction:: glorot_normal

docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.initializers.glorot_uniform.rst

@@ -0,0 +1,6 @@
+mlx.nn.initializers.glorot\_uniform
+===================================
+
+.. currentmodule:: mlx.nn.initializers
+
+.. autofunction:: glorot_uniform

docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.initializers.he_normal.rst

@@ -0,0 +1,6 @@
+mlx.nn.initializers.he\_normal
+==============================
+
+.. currentmodule:: mlx.nn.initializers
+
+.. autofunction:: he_normal

docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.initializers.he_uniform.rst

@@ -0,0 +1,6 @@
+mlx.nn.initializers.he\_uniform
+===============================
+
+.. currentmodule:: mlx.nn.initializers
+
+.. autofunction:: he_uniform

docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.initializers.identity.rst

@@ -0,0 +1,6 @@
+mlx.nn.initializers.identity
+============================
+
+.. currentmodule:: mlx.nn.initializers
+
+.. autofunction:: identity

docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.initializers.normal.rst

@@ -0,0 +1,6 @@
+mlx.nn.initializers.normal
+==========================
+
+.. currentmodule:: mlx.nn.initializers
+
+.. autofunction:: normal

docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.initializers.uniform.rst

@@ -0,0 +1,6 @@
+mlx.nn.initializers.uniform
+===========================
+
+.. currentmodule:: mlx.nn.initializers
+
+.. autofunction:: uniform

docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.losses.gaussian_nll_loss.rst

@@ -0,0 +1,8 @@
+mlx.nn.losses.gaussian\_nll\_loss
+=================================
+
+.. currentmodule:: mlx.nn.losses
+
+.. autoclass:: gaussian_nll_loss
+
+   

docs/build/html/_sources/python/nn/_autosummary_functions/mlx.nn.softshrink.rst

@@ -0,0 +1,8 @@
+mlx.nn.softshrink
+=================
+
+.. currentmodule:: mlx.nn
+
+.. autoclass:: softshrink
+
+   

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

@@ -19,5 +19,6 @@ simple functions.
    prelu
    relu
    selu
+   softshrink
    silu
    step

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

@@ -0,0 +1,45 @@
+.. _init:
+
+.. currentmodule:: mlx.nn.init
+
+Initializers
+------------
+
+The ``mlx.nn.init`` package contains commonly used initializers for neural
+network parameters. Initializers return a function which can be applied to any
+input :obj:`mlx.core.array` to produce an initialized output.
+
+For example:
+
+.. code:: python
+
+   import mlx.core as mx
+   import mlx.nn as nn
+
+   init_fn = nn.init.uniform()
+
+   # Produces a [2, 2] uniform matrix
+   param = init_fn(mx.zeros((2, 2)))
+
+To re-initialize all the parameter in an :obj:`mlx.nn.Module` from say a uniform 
+distribution, you can do:
+
+.. code:: python
+  
+   import mlx.nn as nn
+   model = nn.Sequential(nn.Linear(5, 10), nn.ReLU(), nn.Linear(10, 5))
+   init_fn = nn.init.uniform(low=-0.1, high=0.1)
+   model.apply(init_fn)
+   
+
+.. autosummary::
+   :toctree: _autosummary
+
+   constant
+   normal
+   uniform
+   identity
+   glorot_normal
+   glorot_uniform
+   he_normal
+   he_uniform

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

@@ -0,0 +1,18 @@
+.. _initializers:
+
+.. currentmodule:: mlx.nn.initializers
+
+Initializers
+--------------
+
+.. autosummary::
+   :toctree: _autosummary_functions
+
+   constant
+   normal
+   uniform
+   identity
+   glorot_normal
+   glorot_uniform
+   he_normal
+   he_uniform

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

@@ -33,5 +33,6 @@ Layers
    Sequential
    SiLU
    SinusoidalPositionalEncoding
+   Softshrink
    Step
    Transformer

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

@@ -12,6 +12,7 @@ Loss Functions
    binary_cross_entropy
    cosine_similarity_loss
    cross_entropy
+   gaussian_nll_loss
    hinge_loss
    huber_loss
    kl_div_loss