docs update

docs/build/html/_sources/python/_autosummary/mlx.core.custom_function.jvp.rst

@@ -1,6 +0,0 @@
-mlx.core.custom\_function.jvp
-=============================
-
-.. currentmodule:: mlx.core
-
-.. automethod:: custom_function.jvp

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

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

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

@@ -0,0 +1,6 @@
+mlx.core.einsum\_path
+=====================
+
+.. currentmodule:: mlx.core
+
+.. autofunction:: einsum_path

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

@@ -0,0 +1,6 @@
+mlx.core.nan\_to\_num
+=====================
+
+.. currentmodule:: mlx.core
+
+.. autofunction:: nan_to_num

docs/build/html/_sources/python/_autosummary/mlx.core.random.laplace.rst

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

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

@@ -57,6 +57,8 @@ Operations
    diagonal
    divide
    divmod
+   einsum
+   einsum_path
    equal
    erf
    erfinv
@@ -104,6 +106,7 @@ Operations
    minimum
    moveaxis
    multiply
+   nan_to_num
    negative
    not_equal
    ones

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

@@ -31,6 +31,41 @@ model's parameters and the **optimizer state**.
             # Compute the new parameters but also the optimizer state.
             mx.eval(model.parameters(), optimizer.state)
 
+Saving and Loading
+------------------
+
+To serialize an optimizer, save its state. To load an optimizer, load and set
+the saved state. Here's a simple example:
+
+.. code-block:: python
+
+   import mlx.core as mx
+   from mlx.utils import tree_flatten, tree_unflatten
+   import mlx.optimizers as optim
+
+   optimizer = optim.Adam(learning_rate=1e-2)
+
+   # Perform some updates with the optimizer
+   model = {"w" : mx.zeros((5, 5))}
+   grads = {"w" : mx.ones((5, 5))}
+   optimizer.update(model, grads)
+
+   # Save the state
+   state = tree_flatten(optimizer.state)
+   mx.save_safetensors("optimizer.safetensors", dict(state))
+
+   # Later on, for example when loading from a checkpoint,
+   # recreate the optimizer and load the state
+   optimizer = optim.Adam(learning_rate=1e-2)
+
+   state = tree_unflatten(list(mx.load("optimizer.safetensors").items()))
+   optimizer.state = state
+
+Note, not every optimizer configuation parameter is saved in the state. For
+example, for Adam the learning rate is saved but the ``betas`` and ``eps``
+parameters are not. A good rule of thumb is if the parameter can be scheduled
+then it will be included in the optimizer state.
+
 .. toctree::
 
    optimizers/optimizer

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

@@ -44,3 +44,4 @@ we use a splittable version of Threefry, which is a counter-based PRNG.
    split
    truncated_normal
    uniform
+   laplace