docs update

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

@@ -63,6 +63,8 @@
    
       ~array.T
       ~array.dtype
+      ~array.itemsize
+      ~array.nbytes
       ~array.ndim
       ~array.shape
       ~array.size

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

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

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

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

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

@@ -0,0 +1,6 @@
+mlx.core.save\_safetensors
+==========================
+
+.. currentmodule:: mlx.core
+
+.. autofunction:: save_safetensors

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

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

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

@@ -0,0 +1,11 @@
+.. _linalg:
+
+Linear Algebra
+==============
+
+.. currentmodule:: mlx.core.linalg
+
+.. autosummary:: 
+   :toctree: _autosummary 
+
+    norm

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

@@ -123,7 +123,7 @@ To get more detailed information on the arrays in a :class:`Module` you can use
 all the parameters in a :class:`Module` do:
 
 .. code-block:: python
-    
+
    from mlx.utils import tree_map
    shapes = tree_map(lambda p: p.shape, mlp.parameters())
 
@@ -131,7 +131,7 @@ As another example, you can count the number of parameters in a :class:`Module`
 with:
 
 .. code-block:: python
-    
+
    from mlx.utils import tree_flatten
    num_params = sum(v.size for _, v in tree_flatten(mlp.parameters()))
 
@@ -170,14 +170,13 @@ In detail:
   :meth:`mlx.core.value_and_grad`
 
 .. autosummary::
-   :recursive:
    :toctree: _autosummary
 
    value_and_grad
-   Module
 
 .. toctree::
 
+   nn/module
    nn/layers
    nn/functions
    nn/losses

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

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

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

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

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

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

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

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

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

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

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

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

docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Module.apply.rst

@@ -0,0 +1,6 @@
+mlx.nn.Module.apply
+===================
+
+.. currentmodule:: mlx.nn
+
+.. automethod:: Module.apply

docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Module.apply_to_modules.rst

@@ -0,0 +1,6 @@
+mlx.nn.Module.apply\_to\_modules
+================================
+
+.. currentmodule:: mlx.nn
+
+.. automethod:: Module.apply_to_modules

docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Module.children.rst

@@ -0,0 +1,6 @@
+mlx.nn.Module.children
+======================
+
+.. currentmodule:: mlx.nn
+
+.. automethod:: Module.children

docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Module.eval.rst

@@ -0,0 +1,6 @@
+mlx.nn.Module.eval
+==================
+
+.. currentmodule:: mlx.nn
+
+.. automethod:: Module.eval

docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Module.filter_and_map.rst

@@ -0,0 +1,6 @@
+mlx.nn.Module.filter\_and\_map
+==============================
+
+.. currentmodule:: mlx.nn
+
+.. automethod:: Module.filter_and_map

docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Module.freeze.rst

@@ -0,0 +1,6 @@
+mlx.nn.Module.freeze
+====================
+
+.. currentmodule:: mlx.nn
+
+.. automethod:: Module.freeze

docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Module.leaf_modules.rst

@@ -0,0 +1,6 @@
+mlx.nn.Module.leaf\_modules
+===========================
+
+.. currentmodule:: mlx.nn
+
+.. automethod:: Module.leaf_modules

docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Module.load_weights.rst

@@ -0,0 +1,6 @@
+mlx.nn.Module.load\_weights
+===========================
+
+.. currentmodule:: mlx.nn
+
+.. automethod:: Module.load_weights

docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Module.modules.rst

@@ -0,0 +1,6 @@
+mlx.nn.Module.modules
+=====================
+
+.. currentmodule:: mlx.nn
+
+.. automethod:: Module.modules

docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Module.named_modules.rst

@@ -0,0 +1,6 @@
+mlx.nn.Module.named\_modules
+============================
+
+.. currentmodule:: mlx.nn
+
+.. automethod:: Module.named_modules

docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Module.parameters.rst

@@ -0,0 +1,6 @@
+mlx.nn.Module.parameters
+========================
+
+.. currentmodule:: mlx.nn
+
+.. automethod:: Module.parameters

docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Module.save_weights.rst

@@ -0,0 +1,6 @@
+mlx.nn.Module.save\_weights
+===========================
+
+.. currentmodule:: mlx.nn
+
+.. automethod:: Module.save_weights

docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Module.train.rst

@@ -0,0 +1,6 @@
+mlx.nn.Module.train
+===================
+
+.. currentmodule:: mlx.nn
+
+.. automethod:: Module.train

docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Module.trainable_parameters.rst

@@ -0,0 +1,6 @@
+mlx.nn.Module.trainable\_parameters
+===================================
+
+.. currentmodule:: mlx.nn
+
+.. automethod:: Module.trainable_parameters

docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Module.training.rst

@@ -0,0 +1,6 @@
+mlx.nn.Module.training
+======================
+
+.. currentmodule:: mlx.nn
+
+.. autoproperty:: Module.training

docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Module.unfreeze.rst

@@ -0,0 +1,6 @@
+mlx.nn.Module.unfreeze
+======================
+
+.. currentmodule:: mlx.nn
+
+.. automethod:: Module.unfreeze

docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Module.update.rst

@@ -0,0 +1,6 @@
+mlx.nn.Module.update
+====================
+
+.. currentmodule:: mlx.nn
+
+.. automethod:: Module.update

docs/build/html/_sources/python/nn/_autosummary/mlx.nn.Module.update_modules.rst

@@ -0,0 +1,6 @@
+mlx.nn.Module.update\_modules
+=============================
+
+.. currentmodule:: mlx.nn
+
+.. automethod:: Module.update_modules

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

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

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

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

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

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

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

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

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

@@ -0,0 +1,8 @@
+mlx.nn.losses.log\_cosh\_loss
+=============================
+
+.. currentmodule:: mlx.nn.losses
+
+.. autoclass:: log_cosh_loss
+
+   

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

@@ -9,7 +9,7 @@ Layers
    :toctree: _autosummary
    :template: nn-module-template.rst
 
-   Embedding
+   Sequential
    ReLU
    PReLU
    GELU
@@ -17,13 +17,21 @@ Layers
    Step
    SELU
    Mish
+   Embedding
    Linear
+   QuantizedLinear
    Conv1d
    Conv2d
+   BatchNorm
    LayerNorm
    RMSNorm
    GroupNorm
-   RoPE
+   InstanceNorm
+   Dropout
+   Dropout2d
+   Dropout3d
+   Transformer
    MultiHeadAttention
-   Sequential
-   QuantizedLinear
+   ALiBi
+   RoPE
+   SinusoidalPositionalEncoding

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

@@ -16,4 +16,7 @@ Loss Functions
    mse_loss
    nll_loss
    smooth_l1_loss
-   triplet_loss
+   triplet_loss
+   hinge_loss
+   huber_loss
+   log_cosh_loss

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

@@ -0,0 +1,36 @@
+Module
+======
+
+.. currentmodule:: mlx.nn
+
+.. autoclass:: Module
+
+   .. rubric:: Attributes
+
+   .. autosummary::
+      :toctree: _autosummary
+   
+      Module.training
+   
+   .. rubric:: Methods
+
+   .. autosummary::
+      :toctree: _autosummary
+   
+      Module.apply
+      Module.apply_to_modules
+      Module.children
+      Module.eval
+      Module.filter_and_map
+      Module.freeze
+      Module.leaf_modules
+      Module.load_weights
+      Module.modules
+      Module.named_modules
+      Module.parameters
+      Module.save_weights
+      Module.train
+      Module.trainable_parameters
+      Module.unfreeze
+      Module.update
+      Module.update_modules

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

@@ -77,12 +77,14 @@ Operations
    quantize
    quantized_matmul
    reciprocal
+   repeat
    reshape
    round
    rsqrt
    save
    savez
    savez_compressed
+   save_safetensors
    sigmoid
    sign
    sin
@@ -102,6 +104,7 @@ Operations
    take_along_axis
    tan
    tanh
+   tensordot
    transpose
    tri
    tril