update docs

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

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

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

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

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

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

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

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

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

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

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

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

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

@@ -0,0 +1,8 @@
+mlx.nn.losses.binary\_cross\_entropy
+====================================
+
+.. currentmodule:: mlx.nn.losses
+
+.. autoclass:: binary_cross_entropy
+
+   

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

@@ -0,0 +1,8 @@
+mlx.nn.losses.cross\_entropy
+============================
+
+.. currentmodule:: mlx.nn.losses
+
+.. autoclass:: cross_entropy
+
+   

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

@@ -0,0 +1,8 @@
+mlx.nn.losses.kl\_div\_loss
+===========================
+
+.. currentmodule:: mlx.nn.losses
+
+.. autoclass:: kl_div_loss
+
+   

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

@@ -0,0 +1,8 @@
+mlx.nn.losses.l1\_loss
+======================
+
+.. currentmodule:: mlx.nn.losses
+
+.. autoclass:: l1_loss
+
+   

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

@@ -0,0 +1,8 @@
+mlx.nn.losses.mse\_loss
+=======================
+
+.. currentmodule:: mlx.nn.losses
+
+.. autoclass:: mse_loss
+
+   

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

@@ -0,0 +1,8 @@
+mlx.nn.losses.nll\_loss
+=======================
+
+.. currentmodule:: mlx.nn.losses
+
+.. autoclass:: nll_loss
+
+   

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

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

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

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

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

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

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

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

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

@@ -97,7 +97,7 @@ Updating the parameters
 
 MLX modules allow accessing and updating individual parameters. However, most
 times we need to update large subsets of a module's parameters. This action is
-performed by :meth:`Module.update`. 
+performed by :meth:`Module.update`.
 
 Value and grad
 --------------
@@ -146,8 +146,12 @@ Neural Network Layers
 
    Embedding
    ReLU
+   PReLU
    GELU
    SiLU
+   Step
+   SELU
+   Mish
    Linear
    Conv1d
    Conv2d
@@ -169,4 +173,22 @@ simple functions.
    gelu_approx
    gelu_fast_approx
    relu
+   prelu
    silu
+   step
+   selu
+   mish
+
+Loss Functions
+--------------
+
+.. autosummary::
+   :toctree: _autosummary_functions
+   :template: nn-module-template.rst
+
+   losses.cross_entropy
+   losses.binary_cross_entropy
+   losses.l1_loss
+   losses.mse_loss
+   losses.nll_loss
+   losses.kl_div_loss

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

@@ -38,9 +38,11 @@ Operations
    erfinv
    exp
    expand_dims
+   eye
    full
    greater
    greater_equal
+   identity
    less
    less_equal
    load

docs/build/html/_sources/unified_memory.rst

@@ -6,7 +6,7 @@ Unified Memory
 .. currentmodule:: mlx.core
 
 Apple silicon has a unified memory architecture. The CPU and GPU have direct
-access to the same memory pool. MLX is designed to take advantage that.
+access to the same memory pool. MLX is designed to take advantage of that.
 
 Concretely, when you make an array in MLX you don't have to specify its location: