update docs

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

@@ -929,7 +929,7 @@ We see some modest improvements right away!
 
 This operation is now good to be used to build other operations, 
 in :class:`mlx.nn.Module` calls, and also as a part of graph 
-transformations such as :meth:`grad` and :meth:`simplify`!
+transformations like :meth:`grad`!
 
 Scripts
 -------

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

@@ -25,6 +25,8 @@
       ~array.cummin
       ~array.cumprod
       ~array.cumsum
+      ~array.diag
+      ~array.diagonal
       ~array.exp
       ~array.flatten
       ~array.item

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

@@ -0,0 +1,6 @@
+mlx.core.diag
+=============
+
+.. currentmodule:: mlx.core
+
+.. autofunction:: diag

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

@@ -0,0 +1,6 @@
+mlx.core.diagonal
+=================
+
+.. currentmodule:: mlx.core
+
+.. autofunction:: diagonal

docs/build/html/_sources/python/_autosummary/mlx.core.linalg.qr.rst

@@ -0,0 +1,6 @@
+mlx.core.linalg.qr
+==================
+
+.. currentmodule:: mlx.core.linalg
+
+.. autofunction:: qr

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

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

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

@@ -9,3 +9,4 @@ Linear Algebra
    :toctree: _autosummary 
 
     norm
+    qr

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

@@ -180,3 +180,4 @@ In detail:
    nn/layers
    nn/functions
    nn/losses
+   nn/init

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

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

@@ -35,6 +35,8 @@ Operations
    cos
    cosh
    dequantize
+   diag
+   diagonal
    divide
    divmod
    equal

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

@@ -40,6 +40,7 @@ model's parameters and the **optimizer state**.
    SGD
    RMSprop
    Adagrad
+   Adafactor
    AdaDelta
    Adam
    AdamW

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

@@ -14,4 +14,3 @@ Transforms
    jvp
    vjp
    vmap
-   simplify

docs/build/html/_sources/usage/lazy_evaluation.rst

@@ -20,7 +20,7 @@ Transforming Compute Graphs
 
 Lazy evaluation let's us record a compute graph without actually doing any
 computations. This is useful for function transformations like :func:`grad` and
-:func:`vmap` and graph optimizations like :func:`simplify`.
+:func:`vmap` and graph optimizations.
 
 Currently, MLX does not compile and rerun compute graphs. They are all
 generated dynamically. However, lazy evaluation makes it much easier to