docs update

docs/build/html/python/optimizers.html

@@ -8,7 +8,7 @@
     <meta charset="utf-8" />
     <meta name="viewport" content="width=device-width, initial-scale=1.0" /><meta name="generator" content="Docutils 0.18.1: http://docutils.sourceforge.net/" />
 
-    <title>Optimizers &#8212; MLX 0.16.0 documentation</title>
+    <title>Optimizers &#8212; MLX 0.16.1 documentation</title>
   
   
   
@@ -36,7 +36,7 @@
 <link rel="preload" as="script" href="../_static/scripts/pydata-sphinx-theme.js?digest=5b4479735964841361fd" />
   <script src="../_static/vendor/fontawesome/6.1.2/js/all.min.js?digest=5b4479735964841361fd"></script>
 
-    <script src="../_static/documentation_options.js?v=ac69e7c2"></script>
+    <script src="../_static/documentation_options.js?v=179565df"></script>
     <script src="../_static/doctools.js?v=888ff710"></script>
     <script src="../_static/sphinx_highlight.js?v=dc90522c"></script>
     <script src="../_static/scripts/sphinx-book-theme.js?v=efea14e4"></script>
@@ -131,8 +131,8 @@
       
     
     
-    <img src="../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.16.0 documentation - Home"/>
-    <script>document.write(`<img src="../_static/mlx_logo_dark.png" class="logo__image only-dark" alt="MLX 0.16.0 documentation - Home"/>`);</script>
+    <img src="../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.16.1 documentation - Home"/>
+    <script>document.write(`<img src="../_static/mlx_logo_dark.png" class="logo__image only-dark" alt="MLX 0.16.1 documentation - Home"/>`);</script>
   
   
 </a></div>
@@ -296,6 +296,8 @@
 <li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.diagonal.html">mlx.core.diagonal</a></li>
 <li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.divide.html">mlx.core.divide</a></li>
 <li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.divmod.html">mlx.core.divmod</a></li>
+<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.einsum.html">mlx.core.einsum</a></li>
+<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.einsum_path.html">mlx.core.einsum_path</a></li>
 <li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.equal.html">mlx.core.equal</a></li>
 <li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.erf.html">mlx.core.erf</a></li>
 <li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.erfinv.html">mlx.core.erfinv</a></li>
@@ -343,6 +345,7 @@
 <li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.minimum.html">mlx.core.minimum</a></li>
 <li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.moveaxis.html">mlx.core.moveaxis</a></li>
 <li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.multiply.html">mlx.core.multiply</a></li>
+<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.nan_to_num.html">mlx.core.nan_to_num</a></li>
 <li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.negative.html">mlx.core.negative</a></li>
 <li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.not_equal.html">mlx.core.not_equal</a></li>
 <li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.ones.html">mlx.core.ones</a></li>
@@ -414,6 +417,7 @@
 <li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.random.split.html">mlx.core.random.split</a></li>
 <li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.random.truncated_normal.html">mlx.core.random.truncated_normal</a></li>
 <li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.random.uniform.html">mlx.core.random.uniform</a></li>
+<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.random.laplace.html">mlx.core.random.laplace</a></li>
 </ul>
 </li>
 <li class="toctree-l1 has-children"><a class="reference internal" href="transforms.html">Transforms</a><input class="toctree-checkbox" id="toctree-checkbox-6" name="toctree-checkbox-6" type="checkbox"/><label class="toctree-toggle" for="toctree-checkbox-6"><i class="fa-solid fa-chevron-down"></i></label><ul>
@@ -803,7 +807,9 @@ document.write(`
   </button>
 `);
 </script>
-
+<label class="sidebar-toggle secondary-toggle btn btn-sm" for="__secondary"title="Toggle secondary sidebar" data-bs-placement="bottom" data-bs-toggle="tooltip">
+    <span class="fa-solid fa-list"></span>
+</label>
 </div></div>
       
     </div>
@@ -819,6 +825,16 @@ document.write(`
     <div id="print-main-content">
         <div id="jb-print-toc">
             
+            <div>
+                <h2> Contents </h2>
+            </div>
+            <nav aria-label="Page">
+                <ul class="visible nav section-nav flex-column">
+<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#saving-and-loading">Saving and Loading</a><ul class="visible nav section-nav flex-column">
+</ul>
+</li>
+</ul>
+            </nav>
         </div>
     </div>
 </div>
@@ -854,6 +870,37 @@ model’s parameters and the <strong>optimizer state</strong>.</p>
         <span class="n">mx</span><span class="o">.</span><span class="n">eval</span><span class="p">(</span><span class="n">model</span><span class="o">.</span><span class="n">parameters</span><span class="p">(),</span> <span class="n">optimizer</span><span class="o">.</span><span class="n">state</span><span class="p">)</span>
 </pre></div>
 </div>
+<section id="saving-and-loading">
+<h2>Saving and Loading<a class="headerlink" href="#saving-and-loading" title="Link to this heading">#</a></h2>
+<p>To serialize an optimizer, save its state. To load an optimizer, load and set
+the saved state. Here’s a simple example:</p>
+<div class="highlight-python notranslate"><div class="highlight"><pre><span></span><span class="kn">import</span> <span class="nn">mlx.core</span> <span class="k">as</span> <span class="nn">mx</span>
+<span class="kn">from</span> <span class="nn">mlx.utils</span> <span class="kn">import</span> <span class="n">tree_flatten</span><span class="p">,</span> <span class="n">tree_unflatten</span>
+<span class="kn">import</span> <span class="nn">mlx.optimizers</span> <span class="k">as</span> <span class="nn">optim</span>
+
+<span class="n">optimizer</span> <span class="o">=</span> <span class="n">optim</span><span class="o">.</span><span class="n">Adam</span><span class="p">(</span><span class="n">learning_rate</span><span class="o">=</span><span class="mf">1e-2</span><span class="p">)</span>
+
+<span class="c1"># Perform some updates with the optimizer</span>
+<span class="n">model</span> <span class="o">=</span> <span class="p">{</span><span class="s2">&quot;w&quot;</span> <span class="p">:</span> <span class="n">mx</span><span class="o">.</span><span class="n">zeros</span><span class="p">((</span><span class="mi">5</span><span class="p">,</span> <span class="mi">5</span><span class="p">))}</span>
+<span class="n">grads</span> <span class="o">=</span> <span class="p">{</span><span class="s2">&quot;w&quot;</span> <span class="p">:</span> <span class="n">mx</span><span class="o">.</span><span class="n">ones</span><span class="p">((</span><span class="mi">5</span><span class="p">,</span> <span class="mi">5</span><span class="p">))}</span>
+<span class="n">optimizer</span><span class="o">.</span><span class="n">update</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">grads</span><span class="p">)</span>
+
+<span class="c1"># Save the state</span>
+<span class="n">state</span> <span class="o">=</span> <span class="n">tree_flatten</span><span class="p">(</span><span class="n">optimizer</span><span class="o">.</span><span class="n">state</span><span class="p">)</span>
+<span class="n">mx</span><span class="o">.</span><span class="n">save_safetensors</span><span class="p">(</span><span class="s2">&quot;optimizer.safetensors&quot;</span><span class="p">,</span> <span class="nb">dict</span><span class="p">(</span><span class="n">state</span><span class="p">))</span>
+
+<span class="c1"># Later on, for example when loading from a checkpoint,</span>
+<span class="c1"># recreate the optimizer and load the state</span>
+<span class="n">optimizer</span> <span class="o">=</span> <span class="n">optim</span><span class="o">.</span><span class="n">Adam</span><span class="p">(</span><span class="n">learning_rate</span><span class="o">=</span><span class="mf">1e-2</span><span class="p">)</span>
+
+<span class="n">state</span> <span class="o">=</span> <span class="n">tree_unflatten</span><span class="p">(</span><span class="nb">list</span><span class="p">(</span><span class="n">mx</span><span class="o">.</span><span class="n">load</span><span class="p">(</span><span class="s2">&quot;optimizer.safetensors&quot;</span><span class="p">)</span><span class="o">.</span><span class="n">items</span><span class="p">()))</span>
+<span class="n">optimizer</span><span class="o">.</span><span class="n">state</span> <span class="o">=</span> <span class="n">state</span>
+</pre></div>
+</div>
+<p>Note, not every optimizer configuation parameter is saved in the state. For
+example, for Adam the learning rate is saved but the <code class="docutils literal notranslate"><span class="pre">betas</span></code> and <code class="docutils literal notranslate"><span class="pre">eps</span></code>
+parameters are not. A good rule of thumb is if the parameter can be scheduled
+then it will be included in the optimizer state.</p>
 <div class="toctree-wrapper compound">
 <ul>
 <li class="toctree-l1"><a class="reference internal" href="optimizers/optimizer.html">Optimizer</a><ul>
@@ -947,6 +994,7 @@ model’s parameters and the <strong>optimizer state</strong>.</p>
 </tr>
 </tbody>
 </table>
+</section>
 </section>
 
 
@@ -985,6 +1033,22 @@ model’s parameters and the <strong>optimizer state</strong>.</p>
             
             
               
+                <div class="bd-sidebar-secondary bd-toc"><div class="sidebar-secondary-items sidebar-secondary__inner">
+
+  <div class="sidebar-secondary-item">
+  <div class="page-toc tocsection onthispage">
+    <i class="fa-solid fa-list"></i> Contents
+  </div>
+  <nav class="bd-toc-nav page-toc">
+    <ul class="visible nav section-nav flex-column">
+<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#saving-and-loading">Saving and Loading</a><ul class="visible nav section-nav flex-column">
+</ul>
+</li>
+</ul>
+  </nav></div>
+
+</div></div>
+              
             
           </div>
           <footer class="bd-footer-content">