diff --git a/docs/build/html/.buildinfo b/docs/build/html/.buildinfo
index 5b5e9de51..63c9d0e41 100644
--- a/docs/build/html/.buildinfo
+++ b/docs/build/html/.buildinfo
@@ -1,4 +1,4 @@
 # Sphinx build info version 1
 # This file hashes the configuration used when building these files. When it is not found, a full rebuild will be done.
-config: 101437d3a6497705badb11e599f96b45
+config: 177bbfa0562989b62a154f2dd1ab02a1
 tags: 645f666f9bcd5a90fca523b33c5a78b7
diff --git a/docs/build/html/_sources/install.rst b/docs/build/html/_sources/install.rst
index b82b2838c..682f09f38 100644
--- a/docs/build/html/_sources/install.rst
+++ b/docs/build/html/_sources/install.rst
@@ -11,10 +11,34 @@ silicon computer is
 
     pip install mlx
 
+To install from PyPI you must meet the following requirements:
+
+- Using an M series chip (Apple silicon)
+- Using a native Python >= 3.8
+- MacOS >= 13.3
+
 .. note::
     MLX is only available on devices running MacOS >= 13.3 
     It is highly recommended to use MacOS 14 (Sonoma)
 
+Troubleshooting
+^^^^^^^^^^^^^^^
+
+*My OS and Python versions are in the required range but pip still does not find
+a matching distribution.*
+
+Probably you are using a non-native Python. The output of
+
+.. code-block:: shell
+
+  python -c "import platform; print(platform.processor())"
+
+should be ``arm``. If it is ``i386`` (and you have M series machine) then you
+are using a non-native Python. Switch your Python to a native Python. A good
+way to do this is with
+`Conda <https://stackoverflow.com/questions/65415996/how-to-specify-the-architecture-or-platform-for-a-new-conda-environment-apple>`_.
+
+
 Build from source
 -----------------
 
@@ -61,6 +85,7 @@ To make sure the install is working run the tests with:
 
 .. code-block:: shell
 
+  pip install ".[testing]"
   python -m unittest discover python/tests
 
 C++ API
diff --git a/docs/build/html/_static/documentation_options.js b/docs/build/html/_static/documentation_options.js
index 9012603b7..f8f4859e3 100644
--- a/docs/build/html/_static/documentation_options.js
+++ b/docs/build/html/_static/documentation_options.js
@@ -1,6 +1,6 @@
 var DOCUMENTATION_OPTIONS = {
     URL_ROOT: document.getElementById("documentation_options").getAttribute('data-url_root'),
-    VERSION: '0.0.3',
+    VERSION: '0.0.4',
     LANGUAGE: 'en',
     COLLAPSE_INDEX: false,
     BUILDER: 'html',
diff --git a/docs/build/html/cpp/ops.html b/docs/build/html/cpp/ops.html
index 4d97b7a65..2167066d3 100644
--- a/docs/build/html/cpp/ops.html
+++ b/docs/build/html/cpp/ops.html
@@ -9,7 +9,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>Operations &#8212; MLX 0.0.3 documentation</title>
+    <title>Operations &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/dev/extensions.html b/docs/build/html/dev/extensions.html
index cb3a13b80..9a47ae5f3 100644
--- a/docs/build/html/dev/extensions.html
+++ b/docs/build/html/dev/extensions.html
@@ -9,7 +9,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>Developer Documentation &#8212; MLX 0.0.3 documentation</title>
+    <title>Developer Documentation &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -133,8 +133,8 @@
       
     
     
-    <img src="../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/examples/linear_regression.html b/docs/build/html/examples/linear_regression.html
index 932db4c26..a096950bf 100644
--- a/docs/build/html/examples/linear_regression.html
+++ b/docs/build/html/examples/linear_regression.html
@@ -9,7 +9,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>Linear Regression &#8212; MLX 0.0.3 documentation</title>
+    <title>Linear Regression &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/examples/llama-inference.html b/docs/build/html/examples/llama-inference.html
index aada78d30..6658e6034 100644
--- a/docs/build/html/examples/llama-inference.html
+++ b/docs/build/html/examples/llama-inference.html
@@ -9,7 +9,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>LLM inference &#8212; MLX 0.0.3 documentation</title>
+    <title>LLM inference &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/examples/mlp.html b/docs/build/html/examples/mlp.html
index 7bf055a43..5ab1e11f8 100644
--- a/docs/build/html/examples/mlp.html
+++ b/docs/build/html/examples/mlp.html
@@ -9,7 +9,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>Multi-Layer Perceptron &#8212; MLX 0.0.3 documentation</title>
+    <title>Multi-Layer Perceptron &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/genindex.html b/docs/build/html/genindex.html
index bd40f287d..4332ef3da 100644
--- a/docs/build/html/genindex.html
+++ b/docs/build/html/genindex.html
@@ -8,7 +8,7 @@
   <head>
     <meta charset="utf-8" />
     <meta name="viewport" content="width=device-width, initial-scale=1.0" />
-    <title>Index &#8212; MLX 0.0.3 documentation</title>
+    <title>Index &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -131,8 +131,8 @@
       
     
     
-    <img src="_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/index.html b/docs/build/html/index.html
index ef7b2e990..c5927c04e 100644
--- a/docs/build/html/index.html
+++ b/docs/build/html/index.html
@@ -9,7 +9,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>MLX &#8212; MLX 0.0.3 documentation</title>
+    <title>MLX &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -133,8 +133,8 @@
       
     
     
-    <img src="_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/install.html b/docs/build/html/install.html
index 45ed70943..f630b7de6 100644
--- a/docs/build/html/install.html
+++ b/docs/build/html/install.html
@@ -9,7 +9,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>Build and Install &#8212; MLX 0.0.3 documentation</title>
+    <title>Build and Install &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
@@ -536,7 +536,10 @@ document.write(`
             </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="#install-from-pypi">Install from PyPI</a></li>
+<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#install-from-pypi">Install from PyPI</a><ul class="visible nav section-nav flex-column">
+<li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#troubleshooting">Troubleshooting</a></li>
+</ul>
+</li>
 <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#build-from-source">Build from source</a><ul class="visible nav section-nav flex-column">
 <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#build-requirements">Build Requirements</a></li>
 <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#python-api">Python API</a></li>
@@ -563,11 +566,30 @@ silicon computer is</p>
 <div class="highlight-shell notranslate"><div class="highlight"><pre><span></span>pip<span class="w"> </span>install<span class="w"> </span>mlx
 </pre></div>
 </div>
+<p>To install from PyPI you must meet the following requirements:</p>
+<ul class="simple">
+<li><p>Using an M series chip (Apple silicon)</p></li>
+<li><p>Using a native Python &gt;= 3.8</p></li>
+<li><p>MacOS &gt;= 13.3</p></li>
+</ul>
 <div class="admonition note">
 <p class="admonition-title">Note</p>
 <p>MLX is only available on devices running MacOS &gt;= 13.3
 It is highly recommended to use MacOS 14 (Sonoma)</p>
 </div>
+<section id="troubleshooting">
+<h3>Troubleshooting<a class="headerlink" href="#troubleshooting" title="Permalink to this heading">#</a></h3>
+<p><em>My OS and Python versions are in the required range but pip still does not find
+a matching distribution.</em></p>
+<p>Probably you are using a non-native Python. The output of</p>
+<div class="highlight-shell notranslate"><div class="highlight"><pre><span></span>python<span class="w"> </span>-c<span class="w"> </span><span class="s2">&quot;import platform; print(platform.processor())&quot;</span>
+</pre></div>
+</div>
+<p>should be <code class="docutils literal notranslate"><span class="pre">arm</span></code>. If it is <code class="docutils literal notranslate"><span class="pre">i386</span></code> (and you have M series machine) then you
+are using a non-native Python. Switch your Python to a native Python. A good
+way to do this is with
+<a class="reference external" href="https://stackoverflow.com/questions/65415996/how-to-specify-the-architecture-or-platform-for-a-new-conda-environment-apple">Conda</a>.</p>
+</section>
 </section>
 <section id="build-from-source">
 <h2>Build from source<a class="headerlink" href="#build-from-source" title="Permalink to this heading">#</a></h2>
@@ -602,7 +624,8 @@ brew<span class="w"> </span>install<span class="w"> </span>pybind11
 </pre></div>
 </div>
 <p>To make sure the install is working run the tests with:</p>
-<div class="highlight-shell notranslate"><div class="highlight"><pre><span></span>python<span class="w"> </span>-m<span class="w"> </span>unittest<span class="w"> </span>discover<span class="w"> </span>python/tests
+<div class="highlight-shell notranslate"><div class="highlight"><pre><span></span>pip<span class="w"> </span>install<span class="w"> </span><span class="s2">&quot;.[testing]&quot;</span>
+python<span class="w"> </span>-m<span class="w"> </span>unittest<span class="w"> </span>discover<span class="w"> </span>python/tests
 </pre></div>
 </div>
 </section>
@@ -722,7 +745,10 @@ MacOS SDK will be used</p>
   </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="#install-from-pypi">Install from PyPI</a></li>
+<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#install-from-pypi">Install from PyPI</a><ul class="visible nav section-nav flex-column">
+<li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#troubleshooting">Troubleshooting</a></li>
+</ul>
+</li>
 <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#build-from-source">Build from source</a><ul class="visible nav section-nav flex-column">
 <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#build-requirements">Build Requirements</a></li>
 <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#python-api">Python API</a></li>
diff --git a/docs/build/html/objects.inv b/docs/build/html/objects.inv
index 9672d198e..d35893603 100644
Binary files a/docs/build/html/objects.inv and b/docs/build/html/objects.inv differ
diff --git a/docs/build/html/python/_autosummary/mlx.core.Device.html b/docs/build/html/python/_autosummary/mlx.core.Device.html
index 00dab276e..d6672f674 100644
--- a/docs/build/html/python/_autosummary/mlx.core.Device.html
+++ b/docs/build/html/python/_autosummary/mlx.core.Device.html
@@ -9,7 +9,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>mlx.core.Device &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.Device &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.Dtype.html b/docs/build/html/python/_autosummary/mlx.core.Dtype.html
index 372f92822..7cca89763 100644
--- a/docs/build/html/python/_autosummary/mlx.core.Dtype.html
+++ b/docs/build/html/python/_autosummary/mlx.core.Dtype.html
@@ -9,7 +9,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>mlx.core.Dtype &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.Dtype &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.Stream.html b/docs/build/html/python/_autosummary/mlx.core.Stream.html
index 4b3a63af2..3d1366677 100644
--- a/docs/build/html/python/_autosummary/mlx.core.Stream.html
+++ b/docs/build/html/python/_autosummary/mlx.core.Stream.html
@@ -9,7 +9,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>mlx.core.Stream &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.Stream &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.abs.html b/docs/build/html/python/_autosummary/mlx.core.abs.html
index 4450f5609..ef211e93a 100644
--- a/docs/build/html/python/_autosummary/mlx.core.abs.html
+++ b/docs/build/html/python/_autosummary/mlx.core.abs.html
@@ -9,7 +9,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>mlx.core.abs &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.abs &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.add.html b/docs/build/html/python/_autosummary/mlx.core.add.html
index 255e5f94f..015833409 100644
--- a/docs/build/html/python/_autosummary/mlx.core.add.html
+++ b/docs/build/html/python/_autosummary/mlx.core.add.html
@@ -9,7 +9,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>mlx.core.add &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.add &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.all.html b/docs/build/html/python/_autosummary/mlx.core.all.html
index e18f0d4da..6cf49ecf0 100644
--- a/docs/build/html/python/_autosummary/mlx.core.all.html
+++ b/docs/build/html/python/_autosummary/mlx.core.all.html
@@ -9,7 +9,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>mlx.core.all &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.all &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.allclose.html b/docs/build/html/python/_autosummary/mlx.core.allclose.html
index 9b7734ffd..7113a9aec 100644
--- a/docs/build/html/python/_autosummary/mlx.core.allclose.html
+++ b/docs/build/html/python/_autosummary/mlx.core.allclose.html
@@ -9,7 +9,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>mlx.core.allclose &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.allclose &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.any.html b/docs/build/html/python/_autosummary/mlx.core.any.html
index cd3fde61c..2c6e12f76 100644
--- a/docs/build/html/python/_autosummary/mlx.core.any.html
+++ b/docs/build/html/python/_autosummary/mlx.core.any.html
@@ -9,7 +9,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>mlx.core.any &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.any &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.arange.html b/docs/build/html/python/_autosummary/mlx.core.arange.html
index c3c9bbca9..df93033fc 100644
--- a/docs/build/html/python/_autosummary/mlx.core.arange.html
+++ b/docs/build/html/python/_autosummary/mlx.core.arange.html
@@ -9,7 +9,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>mlx.core.arange &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.arange &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.arccos.html b/docs/build/html/python/_autosummary/mlx.core.arccos.html
index 798242740..78abe70d7 100644
--- a/docs/build/html/python/_autosummary/mlx.core.arccos.html
+++ b/docs/build/html/python/_autosummary/mlx.core.arccos.html
@@ -9,7 +9,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>mlx.core.arccos &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.arccos &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.arccosh.html b/docs/build/html/python/_autosummary/mlx.core.arccosh.html
index 00f6eb7d9..b84f771af 100644
--- a/docs/build/html/python/_autosummary/mlx.core.arccosh.html
+++ b/docs/build/html/python/_autosummary/mlx.core.arccosh.html
@@ -9,7 +9,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>mlx.core.arccosh &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.arccosh &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.arcsin.html b/docs/build/html/python/_autosummary/mlx.core.arcsin.html
index f7f33d155..2ce19bac0 100644
--- a/docs/build/html/python/_autosummary/mlx.core.arcsin.html
+++ b/docs/build/html/python/_autosummary/mlx.core.arcsin.html
@@ -9,7 +9,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>mlx.core.arcsin &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.arcsin &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.arcsinh.html b/docs/build/html/python/_autosummary/mlx.core.arcsinh.html
index 95c72ccb1..352758ee8 100644
--- a/docs/build/html/python/_autosummary/mlx.core.arcsinh.html
+++ b/docs/build/html/python/_autosummary/mlx.core.arcsinh.html
@@ -9,7 +9,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>mlx.core.arcsinh &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.arcsinh &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.arctan.html b/docs/build/html/python/_autosummary/mlx.core.arctan.html
index 5b04fdc63..c92ac4c0b 100644
--- a/docs/build/html/python/_autosummary/mlx.core.arctan.html
+++ b/docs/build/html/python/_autosummary/mlx.core.arctan.html
@@ -9,7 +9,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>mlx.core.arctan &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.arctan &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.arctanh.html b/docs/build/html/python/_autosummary/mlx.core.arctanh.html
index 9c6e80507..d1f825987 100644
--- a/docs/build/html/python/_autosummary/mlx.core.arctanh.html
+++ b/docs/build/html/python/_autosummary/mlx.core.arctanh.html
@@ -9,7 +9,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>mlx.core.arctanh &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.arctanh &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.argmax.html b/docs/build/html/python/_autosummary/mlx.core.argmax.html
index 25458b60f..3b26565dd 100644
--- a/docs/build/html/python/_autosummary/mlx.core.argmax.html
+++ b/docs/build/html/python/_autosummary/mlx.core.argmax.html
@@ -9,7 +9,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>mlx.core.argmax &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.argmax &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.argmin.html b/docs/build/html/python/_autosummary/mlx.core.argmin.html
index 7a570d90d..12f476cea 100644
--- a/docs/build/html/python/_autosummary/mlx.core.argmin.html
+++ b/docs/build/html/python/_autosummary/mlx.core.argmin.html
@@ -9,7 +9,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>mlx.core.argmin &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.argmin &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.argpartition.html b/docs/build/html/python/_autosummary/mlx.core.argpartition.html
index 4dd9a162f..c6c412a18 100644
--- a/docs/build/html/python/_autosummary/mlx.core.argpartition.html
+++ b/docs/build/html/python/_autosummary/mlx.core.argpartition.html
@@ -9,7 +9,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>mlx.core.argpartition &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.argpartition &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.argsort.html b/docs/build/html/python/_autosummary/mlx.core.argsort.html
index 083defb9f..2f2ea4d33 100644
--- a/docs/build/html/python/_autosummary/mlx.core.argsort.html
+++ b/docs/build/html/python/_autosummary/mlx.core.argsort.html
@@ -9,7 +9,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>mlx.core.argsort &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.argsort &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.array.T.html b/docs/build/html/python/_autosummary/mlx.core.array.T.html
index ef30dce91..6af96ebbd 100644
--- a/docs/build/html/python/_autosummary/mlx.core.array.T.html
+++ b/docs/build/html/python/_autosummary/mlx.core.array.T.html
@@ -9,7 +9,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>mlx.core.array.T &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.array.T &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.array.abs.html b/docs/build/html/python/_autosummary/mlx.core.array.abs.html
index 1973c94b7..7a3d3f662 100644
--- a/docs/build/html/python/_autosummary/mlx.core.array.abs.html
+++ b/docs/build/html/python/_autosummary/mlx.core.array.abs.html
@@ -9,7 +9,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>mlx.core.array.abs &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.array.abs &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.array.all.html b/docs/build/html/python/_autosummary/mlx.core.array.all.html
index 5e8642587..8c8db96cd 100644
--- a/docs/build/html/python/_autosummary/mlx.core.array.all.html
+++ b/docs/build/html/python/_autosummary/mlx.core.array.all.html
@@ -9,7 +9,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>mlx.core.array.all &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.array.all &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.array.any.html b/docs/build/html/python/_autosummary/mlx.core.array.any.html
index 9c3fabf28..843b72f34 100644
--- a/docs/build/html/python/_autosummary/mlx.core.array.any.html
+++ b/docs/build/html/python/_autosummary/mlx.core.array.any.html
@@ -9,7 +9,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>mlx.core.array.any &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.array.any &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.array.argmax.html b/docs/build/html/python/_autosummary/mlx.core.array.argmax.html
index 671031a03..13722b4a9 100644
--- a/docs/build/html/python/_autosummary/mlx.core.array.argmax.html
+++ b/docs/build/html/python/_autosummary/mlx.core.array.argmax.html
@@ -9,7 +9,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>mlx.core.array.argmax &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.array.argmax &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.array.argmin.html b/docs/build/html/python/_autosummary/mlx.core.array.argmin.html
index a9e9ab599..65d0378ec 100644
--- a/docs/build/html/python/_autosummary/mlx.core.array.argmin.html
+++ b/docs/build/html/python/_autosummary/mlx.core.array.argmin.html
@@ -9,7 +9,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>mlx.core.array.argmin &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.array.argmin &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.array.astype.html b/docs/build/html/python/_autosummary/mlx.core.array.astype.html
index 88e7a4c94..e223763e4 100644
--- a/docs/build/html/python/_autosummary/mlx.core.array.astype.html
+++ b/docs/build/html/python/_autosummary/mlx.core.array.astype.html
@@ -9,7 +9,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>mlx.core.array.astype &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.array.astype &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.array.cos.html b/docs/build/html/python/_autosummary/mlx.core.array.cos.html
index eac4222e3..9c5717f1a 100644
--- a/docs/build/html/python/_autosummary/mlx.core.array.cos.html
+++ b/docs/build/html/python/_autosummary/mlx.core.array.cos.html
@@ -9,7 +9,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>mlx.core.array.cos &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.array.cos &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.array.dtype.html b/docs/build/html/python/_autosummary/mlx.core.array.dtype.html
index 1d7ab931a..85372e23f 100644
--- a/docs/build/html/python/_autosummary/mlx.core.array.dtype.html
+++ b/docs/build/html/python/_autosummary/mlx.core.array.dtype.html
@@ -9,7 +9,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>mlx.core.array.dtype &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.array.dtype &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.array.exp.html b/docs/build/html/python/_autosummary/mlx.core.array.exp.html
index 42b5acbfd..bc3ab34fd 100644
--- a/docs/build/html/python/_autosummary/mlx.core.array.exp.html
+++ b/docs/build/html/python/_autosummary/mlx.core.array.exp.html
@@ -9,7 +9,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>mlx.core.array.exp &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.array.exp &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.array.html b/docs/build/html/python/_autosummary/mlx.core.array.html
index 68cea7c7e..b716fdf80 100644
--- a/docs/build/html/python/_autosummary/mlx.core.array.html
+++ b/docs/build/html/python/_autosummary/mlx.core.array.html
@@ -9,7 +9,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>mlx.core.array &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.array &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.array.item.html b/docs/build/html/python/_autosummary/mlx.core.array.item.html
index 820a9bf2d..98fc86142 100644
--- a/docs/build/html/python/_autosummary/mlx.core.array.item.html
+++ b/docs/build/html/python/_autosummary/mlx.core.array.item.html
@@ -9,7 +9,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>mlx.core.array.item &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.array.item &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.array.log.html b/docs/build/html/python/_autosummary/mlx.core.array.log.html
index e854983fb..d32e51560 100644
--- a/docs/build/html/python/_autosummary/mlx.core.array.log.html
+++ b/docs/build/html/python/_autosummary/mlx.core.array.log.html
@@ -9,7 +9,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>mlx.core.array.log &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.array.log &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.array.log1p.html b/docs/build/html/python/_autosummary/mlx.core.array.log1p.html
index f2fcfb7af..e88882a78 100644
--- a/docs/build/html/python/_autosummary/mlx.core.array.log1p.html
+++ b/docs/build/html/python/_autosummary/mlx.core.array.log1p.html
@@ -9,7 +9,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>mlx.core.array.log1p &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.array.log1p &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.array.logsumexp.html b/docs/build/html/python/_autosummary/mlx.core.array.logsumexp.html
index 6a909ec0b..a05fc0a3b 100644
--- a/docs/build/html/python/_autosummary/mlx.core.array.logsumexp.html
+++ b/docs/build/html/python/_autosummary/mlx.core.array.logsumexp.html
@@ -9,7 +9,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>mlx.core.array.logsumexp &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.array.logsumexp &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.array.max.html b/docs/build/html/python/_autosummary/mlx.core.array.max.html
index 32507ebc7..106b413d6 100644
--- a/docs/build/html/python/_autosummary/mlx.core.array.max.html
+++ b/docs/build/html/python/_autosummary/mlx.core.array.max.html
@@ -9,7 +9,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>mlx.core.array.max &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.array.max &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.array.mean.html b/docs/build/html/python/_autosummary/mlx.core.array.mean.html
index c4caa4d06..abc841d1b 100644
--- a/docs/build/html/python/_autosummary/mlx.core.array.mean.html
+++ b/docs/build/html/python/_autosummary/mlx.core.array.mean.html
@@ -9,7 +9,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>mlx.core.array.mean &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.array.mean &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.array.min.html b/docs/build/html/python/_autosummary/mlx.core.array.min.html
index 9b5544183..a92f4b4e6 100644
--- a/docs/build/html/python/_autosummary/mlx.core.array.min.html
+++ b/docs/build/html/python/_autosummary/mlx.core.array.min.html
@@ -9,7 +9,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>mlx.core.array.min &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.array.min &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.array.ndim.html b/docs/build/html/python/_autosummary/mlx.core.array.ndim.html
index dbaa4b799..4ffb766f0 100644
--- a/docs/build/html/python/_autosummary/mlx.core.array.ndim.html
+++ b/docs/build/html/python/_autosummary/mlx.core.array.ndim.html
@@ -9,7 +9,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>mlx.core.array.ndim &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.array.ndim &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.array.prod.html b/docs/build/html/python/_autosummary/mlx.core.array.prod.html
index bffd2c717..46a6abd3c 100644
--- a/docs/build/html/python/_autosummary/mlx.core.array.prod.html
+++ b/docs/build/html/python/_autosummary/mlx.core.array.prod.html
@@ -9,7 +9,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>mlx.core.array.prod &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.array.prod &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.array.reciprocal.html b/docs/build/html/python/_autosummary/mlx.core.array.reciprocal.html
index 685e132b2..602ac5d94 100644
--- a/docs/build/html/python/_autosummary/mlx.core.array.reciprocal.html
+++ b/docs/build/html/python/_autosummary/mlx.core.array.reciprocal.html
@@ -9,7 +9,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>mlx.core.array.reciprocal &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.array.reciprocal &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.array.reshape.html b/docs/build/html/python/_autosummary/mlx.core.array.reshape.html
index 36b725e1f..01950c2aa 100644
--- a/docs/build/html/python/_autosummary/mlx.core.array.reshape.html
+++ b/docs/build/html/python/_autosummary/mlx.core.array.reshape.html
@@ -9,7 +9,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>mlx.core.array.reshape &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.array.reshape &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.array.rsqrt.html b/docs/build/html/python/_autosummary/mlx.core.array.rsqrt.html
index b8bef5147..faae8a579 100644
--- a/docs/build/html/python/_autosummary/mlx.core.array.rsqrt.html
+++ b/docs/build/html/python/_autosummary/mlx.core.array.rsqrt.html
@@ -9,7 +9,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>mlx.core.array.rsqrt &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.array.rsqrt &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.array.shape.html b/docs/build/html/python/_autosummary/mlx.core.array.shape.html
index 983c4ded2..f97d0b33e 100644
--- a/docs/build/html/python/_autosummary/mlx.core.array.shape.html
+++ b/docs/build/html/python/_autosummary/mlx.core.array.shape.html
@@ -9,7 +9,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>mlx.core.array.shape &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.array.shape &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.array.sin.html b/docs/build/html/python/_autosummary/mlx.core.array.sin.html
index cc6095344..2191fb9e3 100644
--- a/docs/build/html/python/_autosummary/mlx.core.array.sin.html
+++ b/docs/build/html/python/_autosummary/mlx.core.array.sin.html
@@ -9,7 +9,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>mlx.core.array.sin &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.array.sin &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.array.size.html b/docs/build/html/python/_autosummary/mlx.core.array.size.html
index 154359df3..4a10bfab9 100644
--- a/docs/build/html/python/_autosummary/mlx.core.array.size.html
+++ b/docs/build/html/python/_autosummary/mlx.core.array.size.html
@@ -9,7 +9,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>mlx.core.array.size &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.array.size &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.array.split.html b/docs/build/html/python/_autosummary/mlx.core.array.split.html
index eb720df1e..155e7f5fa 100644
--- a/docs/build/html/python/_autosummary/mlx.core.array.split.html
+++ b/docs/build/html/python/_autosummary/mlx.core.array.split.html
@@ -9,7 +9,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>mlx.core.array.split &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.array.split &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.array.sqrt.html b/docs/build/html/python/_autosummary/mlx.core.array.sqrt.html
index b9edb5362..8730d9e71 100644
--- a/docs/build/html/python/_autosummary/mlx.core.array.sqrt.html
+++ b/docs/build/html/python/_autosummary/mlx.core.array.sqrt.html
@@ -9,7 +9,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>mlx.core.array.sqrt &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.array.sqrt &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.array.square.html b/docs/build/html/python/_autosummary/mlx.core.array.square.html
index a24b4df3e..cf1313e15 100644
--- a/docs/build/html/python/_autosummary/mlx.core.array.square.html
+++ b/docs/build/html/python/_autosummary/mlx.core.array.square.html
@@ -9,7 +9,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>mlx.core.array.square &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.array.square &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.array.sum.html b/docs/build/html/python/_autosummary/mlx.core.array.sum.html
index 07508bca1..96c7fa0a8 100644
--- a/docs/build/html/python/_autosummary/mlx.core.array.sum.html
+++ b/docs/build/html/python/_autosummary/mlx.core.array.sum.html
@@ -9,7 +9,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>mlx.core.array.sum &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.array.sum &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.array.tolist.html b/docs/build/html/python/_autosummary/mlx.core.array.tolist.html
index deff824a0..ae22d6943 100644
--- a/docs/build/html/python/_autosummary/mlx.core.array.tolist.html
+++ b/docs/build/html/python/_autosummary/mlx.core.array.tolist.html
@@ -9,7 +9,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>mlx.core.array.tolist &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.array.tolist &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.array.transpose.html b/docs/build/html/python/_autosummary/mlx.core.array.transpose.html
index 7d935924a..3929a6a25 100644
--- a/docs/build/html/python/_autosummary/mlx.core.array.transpose.html
+++ b/docs/build/html/python/_autosummary/mlx.core.array.transpose.html
@@ -9,7 +9,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>mlx.core.array.transpose &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.array.transpose &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.array.var.html b/docs/build/html/python/_autosummary/mlx.core.array.var.html
index 4aba45d4f..b8a2628ec 100644
--- a/docs/build/html/python/_autosummary/mlx.core.array.var.html
+++ b/docs/build/html/python/_autosummary/mlx.core.array.var.html
@@ -9,7 +9,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>mlx.core.array.var &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.array.var &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.array_equal.html b/docs/build/html/python/_autosummary/mlx.core.array_equal.html
index 0f087367b..3fa353bfe 100644
--- a/docs/build/html/python/_autosummary/mlx.core.array_equal.html
+++ b/docs/build/html/python/_autosummary/mlx.core.array_equal.html
@@ -9,7 +9,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>mlx.core.array_equal &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.array_equal &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.broadcast_to.html b/docs/build/html/python/_autosummary/mlx.core.broadcast_to.html
index 52da132cc..6254d0ded 100644
--- a/docs/build/html/python/_autosummary/mlx.core.broadcast_to.html
+++ b/docs/build/html/python/_autosummary/mlx.core.broadcast_to.html
@@ -9,7 +9,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>mlx.core.broadcast_to &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.broadcast_to &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.concatenate.html b/docs/build/html/python/_autosummary/mlx.core.concatenate.html
index f33457ee8..2fedca12f 100644
--- a/docs/build/html/python/_autosummary/mlx.core.concatenate.html
+++ b/docs/build/html/python/_autosummary/mlx.core.concatenate.html
@@ -9,7 +9,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>mlx.core.concatenate &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.concatenate &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.conv1d.html b/docs/build/html/python/_autosummary/mlx.core.conv1d.html
index 8119dee13..c94187220 100644
--- a/docs/build/html/python/_autosummary/mlx.core.conv1d.html
+++ b/docs/build/html/python/_autosummary/mlx.core.conv1d.html
@@ -9,7 +9,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>mlx.core.conv1d &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.conv1d &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.conv2d.html b/docs/build/html/python/_autosummary/mlx.core.conv2d.html
index 8fd6c8c75..cd2c87953 100644
--- a/docs/build/html/python/_autosummary/mlx.core.conv2d.html
+++ b/docs/build/html/python/_autosummary/mlx.core.conv2d.html
@@ -9,7 +9,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>mlx.core.conv2d &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.conv2d &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.convolve.html b/docs/build/html/python/_autosummary/mlx.core.convolve.html
index 20d7575c0..c332f3178 100644
--- a/docs/build/html/python/_autosummary/mlx.core.convolve.html
+++ b/docs/build/html/python/_autosummary/mlx.core.convolve.html
@@ -9,7 +9,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>mlx.core.convolve &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.convolve &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.cos.html b/docs/build/html/python/_autosummary/mlx.core.cos.html
index a033b34d5..d25439d14 100644
--- a/docs/build/html/python/_autosummary/mlx.core.cos.html
+++ b/docs/build/html/python/_autosummary/mlx.core.cos.html
@@ -9,7 +9,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>mlx.core.cos &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.cos &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.cosh.html b/docs/build/html/python/_autosummary/mlx.core.cosh.html
index 923931ac4..c4e80cfd9 100644
--- a/docs/build/html/python/_autosummary/mlx.core.cosh.html
+++ b/docs/build/html/python/_autosummary/mlx.core.cosh.html
@@ -9,7 +9,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>mlx.core.cosh &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.cosh &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.default_device.html b/docs/build/html/python/_autosummary/mlx.core.default_device.html
index 240c82e87..fb0b94ccd 100644
--- a/docs/build/html/python/_autosummary/mlx.core.default_device.html
+++ b/docs/build/html/python/_autosummary/mlx.core.default_device.html
@@ -9,7 +9,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>mlx.core.default_device &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.default_device &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.default_stream.html b/docs/build/html/python/_autosummary/mlx.core.default_stream.html
index 3dbd86a68..35c6dbb20 100644
--- a/docs/build/html/python/_autosummary/mlx.core.default_stream.html
+++ b/docs/build/html/python/_autosummary/mlx.core.default_stream.html
@@ -9,7 +9,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>mlx.core.default_stream &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.default_stream &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.divide.html b/docs/build/html/python/_autosummary/mlx.core.divide.html
index a59fefd3c..552a9eda8 100644
--- a/docs/build/html/python/_autosummary/mlx.core.divide.html
+++ b/docs/build/html/python/_autosummary/mlx.core.divide.html
@@ -9,7 +9,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>mlx.core.divide &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.divide &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.equal.html b/docs/build/html/python/_autosummary/mlx.core.equal.html
index 5f765fd67..035b07c1c 100644
--- a/docs/build/html/python/_autosummary/mlx.core.equal.html
+++ b/docs/build/html/python/_autosummary/mlx.core.equal.html
@@ -9,7 +9,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>mlx.core.equal &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.equal &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.erf.html b/docs/build/html/python/_autosummary/mlx.core.erf.html
index c5f1e6871..00195b6ce 100644
--- a/docs/build/html/python/_autosummary/mlx.core.erf.html
+++ b/docs/build/html/python/_autosummary/mlx.core.erf.html
@@ -9,7 +9,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>mlx.core.erf &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.erf &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -135,8 +135,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.erfinv.html b/docs/build/html/python/_autosummary/mlx.core.erfinv.html
index 1aa117cdd..73d5253a2 100644
--- a/docs/build/html/python/_autosummary/mlx.core.erfinv.html
+++ b/docs/build/html/python/_autosummary/mlx.core.erfinv.html
@@ -9,7 +9,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>mlx.core.erfinv &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.erfinv &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.eval.html b/docs/build/html/python/_autosummary/mlx.core.eval.html
index 0cac9dfb4..5b1a7c8d1 100644
--- a/docs/build/html/python/_autosummary/mlx.core.eval.html
+++ b/docs/build/html/python/_autosummary/mlx.core.eval.html
@@ -9,7 +9,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>mlx.core.eval &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.eval &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.exp.html b/docs/build/html/python/_autosummary/mlx.core.exp.html
index 0e85e8b45..a1dcb6198 100644
--- a/docs/build/html/python/_autosummary/mlx.core.exp.html
+++ b/docs/build/html/python/_autosummary/mlx.core.exp.html
@@ -9,7 +9,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>mlx.core.exp &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.exp &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.expand_dims.html b/docs/build/html/python/_autosummary/mlx.core.expand_dims.html
index 302225777..77b252061 100644
--- a/docs/build/html/python/_autosummary/mlx.core.expand_dims.html
+++ b/docs/build/html/python/_autosummary/mlx.core.expand_dims.html
@@ -9,7 +9,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>mlx.core.expand_dims &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.expand_dims &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.fft.fft.html b/docs/build/html/python/_autosummary/mlx.core.fft.fft.html
index cb85e056d..6157734e1 100644
--- a/docs/build/html/python/_autosummary/mlx.core.fft.fft.html
+++ b/docs/build/html/python/_autosummary/mlx.core.fft.fft.html
@@ -9,7 +9,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>mlx.core.fft.fft &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.fft.fft &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.fft.fft2.html b/docs/build/html/python/_autosummary/mlx.core.fft.fft2.html
index 738bf5cb2..68af42d3a 100644
--- a/docs/build/html/python/_autosummary/mlx.core.fft.fft2.html
+++ b/docs/build/html/python/_autosummary/mlx.core.fft.fft2.html
@@ -9,7 +9,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>mlx.core.fft.fft2 &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.fft.fft2 &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.fft.fftn.html b/docs/build/html/python/_autosummary/mlx.core.fft.fftn.html
index f81e2b7c9..414b957cc 100644
--- a/docs/build/html/python/_autosummary/mlx.core.fft.fftn.html
+++ b/docs/build/html/python/_autosummary/mlx.core.fft.fftn.html
@@ -9,7 +9,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>mlx.core.fft.fftn &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.fft.fftn &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.fft.ifft.html b/docs/build/html/python/_autosummary/mlx.core.fft.ifft.html
index 0c9b19232..15338d35a 100644
--- a/docs/build/html/python/_autosummary/mlx.core.fft.ifft.html
+++ b/docs/build/html/python/_autosummary/mlx.core.fft.ifft.html
@@ -9,7 +9,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>mlx.core.fft.ifft &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.fft.ifft &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.fft.ifft2.html b/docs/build/html/python/_autosummary/mlx.core.fft.ifft2.html
index 9c2a22c3f..bc84b6e62 100644
--- a/docs/build/html/python/_autosummary/mlx.core.fft.ifft2.html
+++ b/docs/build/html/python/_autosummary/mlx.core.fft.ifft2.html
@@ -9,7 +9,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>mlx.core.fft.ifft2 &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.fft.ifft2 &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.fft.ifftn.html b/docs/build/html/python/_autosummary/mlx.core.fft.ifftn.html
index 73ded7492..18442a601 100644
--- a/docs/build/html/python/_autosummary/mlx.core.fft.ifftn.html
+++ b/docs/build/html/python/_autosummary/mlx.core.fft.ifftn.html
@@ -9,7 +9,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>mlx.core.fft.ifftn &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.fft.ifftn &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.fft.irfft.html b/docs/build/html/python/_autosummary/mlx.core.fft.irfft.html
index 1ee94e0c4..2e2145a84 100644
--- a/docs/build/html/python/_autosummary/mlx.core.fft.irfft.html
+++ b/docs/build/html/python/_autosummary/mlx.core.fft.irfft.html
@@ -9,7 +9,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>mlx.core.fft.irfft &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.fft.irfft &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.fft.irfft2.html b/docs/build/html/python/_autosummary/mlx.core.fft.irfft2.html
index 73395e0b0..7f15abf47 100644
--- a/docs/build/html/python/_autosummary/mlx.core.fft.irfft2.html
+++ b/docs/build/html/python/_autosummary/mlx.core.fft.irfft2.html
@@ -9,7 +9,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>mlx.core.fft.irfft2 &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.fft.irfft2 &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.fft.irfftn.html b/docs/build/html/python/_autosummary/mlx.core.fft.irfftn.html
index 26e091577..c5e5e899a 100644
--- a/docs/build/html/python/_autosummary/mlx.core.fft.irfftn.html
+++ b/docs/build/html/python/_autosummary/mlx.core.fft.irfftn.html
@@ -9,7 +9,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>mlx.core.fft.irfftn &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.fft.irfftn &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.fft.rfft.html b/docs/build/html/python/_autosummary/mlx.core.fft.rfft.html
index 8991bf8b8..971352c4a 100644
--- a/docs/build/html/python/_autosummary/mlx.core.fft.rfft.html
+++ b/docs/build/html/python/_autosummary/mlx.core.fft.rfft.html
@@ -9,7 +9,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>mlx.core.fft.rfft &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.fft.rfft &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.fft.rfft2.html b/docs/build/html/python/_autosummary/mlx.core.fft.rfft2.html
index 1a071e4a7..ffa9115ed 100644
--- a/docs/build/html/python/_autosummary/mlx.core.fft.rfft2.html
+++ b/docs/build/html/python/_autosummary/mlx.core.fft.rfft2.html
@@ -9,7 +9,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>mlx.core.fft.rfft2 &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.fft.rfft2 &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.fft.rfftn.html b/docs/build/html/python/_autosummary/mlx.core.fft.rfftn.html
index 8bb155586..807ea310f 100644
--- a/docs/build/html/python/_autosummary/mlx.core.fft.rfftn.html
+++ b/docs/build/html/python/_autosummary/mlx.core.fft.rfftn.html
@@ -9,7 +9,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>mlx.core.fft.rfftn &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.fft.rfftn &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.full.html b/docs/build/html/python/_autosummary/mlx.core.full.html
index 9785aa761..189e8ef44 100644
--- a/docs/build/html/python/_autosummary/mlx.core.full.html
+++ b/docs/build/html/python/_autosummary/mlx.core.full.html
@@ -9,7 +9,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>mlx.core.full &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.full &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.grad.html b/docs/build/html/python/_autosummary/mlx.core.grad.html
index 6c85839c6..6c7af1f86 100644
--- a/docs/build/html/python/_autosummary/mlx.core.grad.html
+++ b/docs/build/html/python/_autosummary/mlx.core.grad.html
@@ -9,7 +9,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>mlx.core.grad &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.grad &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
@@ -542,7 +542,7 @@ document.write(`
 <h1>mlx.core.grad<a class="headerlink" href="#mlx-core-grad" title="Permalink to this heading">#</a></h1>
 <dl class="py function">
 <dt class="sig sig-object py" id="mlx.core.grad">
-<span class="sig-prename descclassname"><span class="pre">mlx.core.</span></span><span class="sig-name descname"><span class="pre">grad</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">fun</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">function</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">argnums</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">Union</span><span class="p"><span class="pre">[</span></span><a class="reference external" href="https://docs.python.org/3/library/functions.html#int" title="(in Python v3.12)"><span class="pre">int</span></a><span class="p"><span class="pre">,</span></span><span class="w"> </span><span class="pre">List</span><span class="p"><span class="pre">[</span></span><a class="reference external" href="https://docs.python.org/3/library/functions.html#int" title="(in Python v3.12)"><span class="pre">int</span></a><span class="p"><span class="pre">]</span></span><span class="p"><span class="pre">]</span></span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">0</span></span></em><span class="sig-paren">)</span> <span class="sig-return"><span class="sig-return-icon">&#x2192;</span> <span class="sig-return-typehint"><span class="pre">cpp_function</span></span></span><a class="headerlink" href="#mlx.core.grad" title="Permalink to this definition">#</a></dt>
+<span class="sig-prename descclassname"><span class="pre">mlx.core.</span></span><span class="sig-name descname"><span class="pre">grad</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">fun</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">function</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">argnums</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">Optional</span><span class="p"><span class="pre">[</span></span><span class="pre">Union</span><span class="p"><span class="pre">[</span></span><a class="reference external" href="https://docs.python.org/3/library/functions.html#int" title="(in Python v3.12)"><span class="pre">int</span></a><span class="p"><span class="pre">,</span></span><span class="w"> </span><span class="pre">List</span><span class="p"><span class="pre">[</span></span><a class="reference external" href="https://docs.python.org/3/library/functions.html#int" title="(in Python v3.12)"><span class="pre">int</span></a><span class="p"><span class="pre">]</span></span><span class="p"><span class="pre">]</span></span><span class="p"><span class="pre">]</span></span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">None</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">argnames</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">Union</span><span class="p"><span class="pre">[</span></span><a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.12)"><span class="pre">str</span></a><span class="p"><span class="pre">,</span></span><span class="w"> </span><span class="pre">List</span><span class="p"><span class="pre">[</span></span><a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.12)"><span class="pre">str</span></a><span class="p"><span class="pre">]</span></span><span class="p"><span class="pre">]</span></span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">[]</span></span></em><span class="sig-paren">)</span> <span class="sig-return"><span class="sig-return-icon">&#x2192;</span> <span class="sig-return-typehint"><span class="pre">cpp_function</span></span></span><a class="headerlink" href="#mlx.core.grad" title="Permalink to this definition">#</a></dt>
 <dd><p>Returns a function which computes the gradient of <code class="docutils literal notranslate"><span class="pre">fun</span></code>.</p>
 <dl class="field-list simple">
 <dt class="field-odd">Parameters<span class="colon">:</span></dt>
@@ -551,8 +551,13 @@ document.write(`
 <a class="reference internal" href="mlx.core.array.html#mlx.core.array" title="mlx.core.array"><code class="xref py py-class docutils literal notranslate"><span class="pre">array</span></code></a> or trees of <a class="reference internal" href="mlx.core.array.html#mlx.core.array" title="mlx.core.array"><code class="xref py py-class docutils literal notranslate"><span class="pre">array</span></code></a> and returns
 a scalar output <a class="reference internal" href="mlx.core.array.html#mlx.core.array" title="mlx.core.array"><code class="xref py py-class docutils literal notranslate"><span class="pre">array</span></code></a>.</p></li>
 <li><p><strong>argnums</strong> (<a class="reference external" href="https://docs.python.org/3/library/functions.html#int" title="(in Python v3.12)"><em>int</em></a><em> or </em><a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#list" title="(in Python v3.12)"><em>list</em></a><em>(</em><a class="reference external" href="https://docs.python.org/3/library/functions.html#int" title="(in Python v3.12)"><em>int</em></a><em>)</em><em>, </em><em>optional</em>) – Specify the index (or indices)
-of the arguments of <code class="docutils literal notranslate"><span class="pre">fun</span></code> to compute the gradient with respect to. If
-unspecified this defaults to <code class="docutils literal notranslate"><span class="pre">0</span></code> indicating <code class="docutils literal notranslate"><span class="pre">fun</span></code>’s first argument.</p></li>
+of the positional arguments of <code class="docutils literal notranslate"><span class="pre">fun</span></code> to compute the gradient
+with respect to. If neither <code class="docutils literal notranslate"><span class="pre">argnums</span></code> nor <code class="docutils literal notranslate"><span class="pre">argnames</span></code> are
+provided <code class="docutils literal notranslate"><span class="pre">argnums</span></code> defaults to <code class="docutils literal notranslate"><span class="pre">0</span></code> indicating <code class="docutils literal notranslate"><span class="pre">fun</span></code>’s first
+argument.</p></li>
+<li><p><strong>argnames</strong> (<a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.12)"><em>str</em></a><em> or </em><a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#list" title="(in Python v3.12)"><em>list</em></a><em>(</em><a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.12)"><em>str</em></a><em>)</em><em>, </em><em>optional</em>) – Specify keyword arguments of
+<code class="docutils literal notranslate"><span class="pre">fun</span></code> to compute gradients with respect to. It defaults to [] so
+no gradients for keyword arguments by default.</p></li>
 </ul>
 </dd>
 <dt class="field-even">Returns<span class="colon">:</span></dt>
diff --git a/docs/build/html/python/_autosummary/mlx.core.greater.html b/docs/build/html/python/_autosummary/mlx.core.greater.html
index 01ccd462c..94ffe5ac4 100644
--- a/docs/build/html/python/_autosummary/mlx.core.greater.html
+++ b/docs/build/html/python/_autosummary/mlx.core.greater.html
@@ -9,7 +9,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>mlx.core.greater &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.greater &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.greater_equal.html b/docs/build/html/python/_autosummary/mlx.core.greater_equal.html
index 884175ab8..f42cf7197 100644
--- a/docs/build/html/python/_autosummary/mlx.core.greater_equal.html
+++ b/docs/build/html/python/_autosummary/mlx.core.greater_equal.html
@@ -9,7 +9,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>mlx.core.greater_equal &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.greater_equal &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.jvp.html b/docs/build/html/python/_autosummary/mlx.core.jvp.html
index a143cba89..1438995aa 100644
--- a/docs/build/html/python/_autosummary/mlx.core.jvp.html
+++ b/docs/build/html/python/_autosummary/mlx.core.jvp.html
@@ -9,7 +9,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>mlx.core.jvp &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.jvp &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.less.html b/docs/build/html/python/_autosummary/mlx.core.less.html
index 433a7c13b..7afd24fc8 100644
--- a/docs/build/html/python/_autosummary/mlx.core.less.html
+++ b/docs/build/html/python/_autosummary/mlx.core.less.html
@@ -9,7 +9,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>mlx.core.less &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.less &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.less_equal.html b/docs/build/html/python/_autosummary/mlx.core.less_equal.html
index b73fda353..e2a20ebdd 100644
--- a/docs/build/html/python/_autosummary/mlx.core.less_equal.html
+++ b/docs/build/html/python/_autosummary/mlx.core.less_equal.html
@@ -9,7 +9,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>mlx.core.less_equal &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.less_equal &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.load.html b/docs/build/html/python/_autosummary/mlx.core.load.html
index f01a61c47..bc99cd0f0 100644
--- a/docs/build/html/python/_autosummary/mlx.core.load.html
+++ b/docs/build/html/python/_autosummary/mlx.core.load.html
@@ -9,7 +9,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>mlx.core.load &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.load &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.log.html b/docs/build/html/python/_autosummary/mlx.core.log.html
index b6635aa94..a4b433a51 100644
--- a/docs/build/html/python/_autosummary/mlx.core.log.html
+++ b/docs/build/html/python/_autosummary/mlx.core.log.html
@@ -9,7 +9,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>mlx.core.log &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.log &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.log10.html b/docs/build/html/python/_autosummary/mlx.core.log10.html
index cafea7dbf..9d209ce1e 100644
--- a/docs/build/html/python/_autosummary/mlx.core.log10.html
+++ b/docs/build/html/python/_autosummary/mlx.core.log10.html
@@ -9,7 +9,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>mlx.core.log10 &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.log10 &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.log1p.html b/docs/build/html/python/_autosummary/mlx.core.log1p.html
index 671eb3ed4..8d43c3d75 100644
--- a/docs/build/html/python/_autosummary/mlx.core.log1p.html
+++ b/docs/build/html/python/_autosummary/mlx.core.log1p.html
@@ -9,7 +9,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>mlx.core.log1p &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.log1p &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.log2.html b/docs/build/html/python/_autosummary/mlx.core.log2.html
index c38c46b2e..f158f5904 100644
--- a/docs/build/html/python/_autosummary/mlx.core.log2.html
+++ b/docs/build/html/python/_autosummary/mlx.core.log2.html
@@ -9,7 +9,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>mlx.core.log2 &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.log2 &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.logaddexp.html b/docs/build/html/python/_autosummary/mlx.core.logaddexp.html
index e467bbe91..239854185 100644
--- a/docs/build/html/python/_autosummary/mlx.core.logaddexp.html
+++ b/docs/build/html/python/_autosummary/mlx.core.logaddexp.html
@@ -9,7 +9,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>mlx.core.logaddexp &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.logaddexp &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.logical_not.html b/docs/build/html/python/_autosummary/mlx.core.logical_not.html
index b2afa8804..b1d3fe7ed 100644
--- a/docs/build/html/python/_autosummary/mlx.core.logical_not.html
+++ b/docs/build/html/python/_autosummary/mlx.core.logical_not.html
@@ -9,7 +9,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>mlx.core.logical_not &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.logical_not &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.logsumexp.html b/docs/build/html/python/_autosummary/mlx.core.logsumexp.html
index f0a2ff1a0..3c2e745f0 100644
--- a/docs/build/html/python/_autosummary/mlx.core.logsumexp.html
+++ b/docs/build/html/python/_autosummary/mlx.core.logsumexp.html
@@ -9,7 +9,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>mlx.core.logsumexp &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.logsumexp &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.matmul.html b/docs/build/html/python/_autosummary/mlx.core.matmul.html
index 0c4603070..b86527cdb 100644
--- a/docs/build/html/python/_autosummary/mlx.core.matmul.html
+++ b/docs/build/html/python/_autosummary/mlx.core.matmul.html
@@ -9,7 +9,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>mlx.core.matmul &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.matmul &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.max.html b/docs/build/html/python/_autosummary/mlx.core.max.html
index d75cc9a5c..82ced27bc 100644
--- a/docs/build/html/python/_autosummary/mlx.core.max.html
+++ b/docs/build/html/python/_autosummary/mlx.core.max.html
@@ -9,7 +9,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>mlx.core.max &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.max &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.maximum.html b/docs/build/html/python/_autosummary/mlx.core.maximum.html
index 73437e28f..a642dfbd2 100644
--- a/docs/build/html/python/_autosummary/mlx.core.maximum.html
+++ b/docs/build/html/python/_autosummary/mlx.core.maximum.html
@@ -9,7 +9,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>mlx.core.maximum &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.maximum &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.mean.html b/docs/build/html/python/_autosummary/mlx.core.mean.html
index defd23e62..484227578 100644
--- a/docs/build/html/python/_autosummary/mlx.core.mean.html
+++ b/docs/build/html/python/_autosummary/mlx.core.mean.html
@@ -9,7 +9,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>mlx.core.mean &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.mean &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.min.html b/docs/build/html/python/_autosummary/mlx.core.min.html
index e3a742a98..ece547627 100644
--- a/docs/build/html/python/_autosummary/mlx.core.min.html
+++ b/docs/build/html/python/_autosummary/mlx.core.min.html
@@ -9,7 +9,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>mlx.core.min &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.min &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.minimum.html b/docs/build/html/python/_autosummary/mlx.core.minimum.html
index 30382d3d0..9777d57a2 100644
--- a/docs/build/html/python/_autosummary/mlx.core.minimum.html
+++ b/docs/build/html/python/_autosummary/mlx.core.minimum.html
@@ -9,7 +9,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>mlx.core.minimum &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.minimum &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.multiply.html b/docs/build/html/python/_autosummary/mlx.core.multiply.html
index c47797912..b9a0e5ea5 100644
--- a/docs/build/html/python/_autosummary/mlx.core.multiply.html
+++ b/docs/build/html/python/_autosummary/mlx.core.multiply.html
@@ -9,7 +9,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>mlx.core.multiply &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.multiply &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.negative.html b/docs/build/html/python/_autosummary/mlx.core.negative.html
index ef05bb1bb..c088cb503 100644
--- a/docs/build/html/python/_autosummary/mlx.core.negative.html
+++ b/docs/build/html/python/_autosummary/mlx.core.negative.html
@@ -9,7 +9,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>mlx.core.negative &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.negative &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.new_stream.html b/docs/build/html/python/_autosummary/mlx.core.new_stream.html
index 1856b770b..0d48d7963 100644
--- a/docs/build/html/python/_autosummary/mlx.core.new_stream.html
+++ b/docs/build/html/python/_autosummary/mlx.core.new_stream.html
@@ -9,7 +9,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>mlx.core.new_stream &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.new_stream &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.ones.html b/docs/build/html/python/_autosummary/mlx.core.ones.html
index 5d37807e3..182e81d11 100644
--- a/docs/build/html/python/_autosummary/mlx.core.ones.html
+++ b/docs/build/html/python/_autosummary/mlx.core.ones.html
@@ -9,7 +9,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>mlx.core.ones &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.ones &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.ones_like.html b/docs/build/html/python/_autosummary/mlx.core.ones_like.html
index 25d07e382..733b442de 100644
--- a/docs/build/html/python/_autosummary/mlx.core.ones_like.html
+++ b/docs/build/html/python/_autosummary/mlx.core.ones_like.html
@@ -9,7 +9,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>mlx.core.ones_like &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.ones_like &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.pad.html b/docs/build/html/python/_autosummary/mlx.core.pad.html
index 85af45d54..70c681566 100644
--- a/docs/build/html/python/_autosummary/mlx.core.pad.html
+++ b/docs/build/html/python/_autosummary/mlx.core.pad.html
@@ -9,7 +9,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>mlx.core.pad &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.pad &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.partition.html b/docs/build/html/python/_autosummary/mlx.core.partition.html
index da66a9582..c7159e0cd 100644
--- a/docs/build/html/python/_autosummary/mlx.core.partition.html
+++ b/docs/build/html/python/_autosummary/mlx.core.partition.html
@@ -9,7 +9,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>mlx.core.partition &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.partition &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.prod.html b/docs/build/html/python/_autosummary/mlx.core.prod.html
index b62e80315..c74e9d5d8 100644
--- a/docs/build/html/python/_autosummary/mlx.core.prod.html
+++ b/docs/build/html/python/_autosummary/mlx.core.prod.html
@@ -9,7 +9,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>mlx.core.prod &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.prod &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.random.bernoulli.html b/docs/build/html/python/_autosummary/mlx.core.random.bernoulli.html
index 3f623bd0b..dfe9596c4 100644
--- a/docs/build/html/python/_autosummary/mlx.core.random.bernoulli.html
+++ b/docs/build/html/python/_autosummary/mlx.core.random.bernoulli.html
@@ -9,7 +9,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>mlx.core.random.bernoulli &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.random.bernoulli &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.random.categorical.html b/docs/build/html/python/_autosummary/mlx.core.random.categorical.html
index f09436432..e54cf9e02 100644
--- a/docs/build/html/python/_autosummary/mlx.core.random.categorical.html
+++ b/docs/build/html/python/_autosummary/mlx.core.random.categorical.html
@@ -9,7 +9,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>mlx.core.random.categorical &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.random.categorical &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.random.gumbel.html b/docs/build/html/python/_autosummary/mlx.core.random.gumbel.html
index 97186361c..ca2894546 100644
--- a/docs/build/html/python/_autosummary/mlx.core.random.gumbel.html
+++ b/docs/build/html/python/_autosummary/mlx.core.random.gumbel.html
@@ -9,7 +9,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>mlx.core.random.gumbel &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.random.gumbel &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.random.key.html b/docs/build/html/python/_autosummary/mlx.core.random.key.html
index 08061453e..38605d0b1 100644
--- a/docs/build/html/python/_autosummary/mlx.core.random.key.html
+++ b/docs/build/html/python/_autosummary/mlx.core.random.key.html
@@ -9,7 +9,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>mlx.core.random.key &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.random.key &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.random.normal.html b/docs/build/html/python/_autosummary/mlx.core.random.normal.html
index 61dff3952..48774355a 100644
--- a/docs/build/html/python/_autosummary/mlx.core.random.normal.html
+++ b/docs/build/html/python/_autosummary/mlx.core.random.normal.html
@@ -9,7 +9,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>mlx.core.random.normal &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.random.normal &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.random.randint.html b/docs/build/html/python/_autosummary/mlx.core.random.randint.html
index c02c5f166..f3b280e0a 100644
--- a/docs/build/html/python/_autosummary/mlx.core.random.randint.html
+++ b/docs/build/html/python/_autosummary/mlx.core.random.randint.html
@@ -9,7 +9,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>mlx.core.random.randint &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.random.randint &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.random.seed.html b/docs/build/html/python/_autosummary/mlx.core.random.seed.html
index d8b71e3ab..3ca52c814 100644
--- a/docs/build/html/python/_autosummary/mlx.core.random.seed.html
+++ b/docs/build/html/python/_autosummary/mlx.core.random.seed.html
@@ -9,7 +9,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>mlx.core.random.seed &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.random.seed &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.random.split.html b/docs/build/html/python/_autosummary/mlx.core.random.split.html
index 567faec12..8a757941c 100644
--- a/docs/build/html/python/_autosummary/mlx.core.random.split.html
+++ b/docs/build/html/python/_autosummary/mlx.core.random.split.html
@@ -9,7 +9,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>mlx.core.random.split &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.random.split &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.random.truncated_normal.html b/docs/build/html/python/_autosummary/mlx.core.random.truncated_normal.html
index 6344ccfb6..16cf0d34c 100644
--- a/docs/build/html/python/_autosummary/mlx.core.random.truncated_normal.html
+++ b/docs/build/html/python/_autosummary/mlx.core.random.truncated_normal.html
@@ -9,7 +9,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>mlx.core.random.truncated_normal &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.random.truncated_normal &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.random.uniform.html b/docs/build/html/python/_autosummary/mlx.core.random.uniform.html
index 4a9f80b3c..8147a75ff 100644
--- a/docs/build/html/python/_autosummary/mlx.core.random.uniform.html
+++ b/docs/build/html/python/_autosummary/mlx.core.random.uniform.html
@@ -9,7 +9,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>mlx.core.random.uniform &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.random.uniform &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.reciprocal.html b/docs/build/html/python/_autosummary/mlx.core.reciprocal.html
index 99ff97ae6..ebfa1c3b9 100644
--- a/docs/build/html/python/_autosummary/mlx.core.reciprocal.html
+++ b/docs/build/html/python/_autosummary/mlx.core.reciprocal.html
@@ -9,7 +9,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>mlx.core.reciprocal &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.reciprocal &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.reshape.html b/docs/build/html/python/_autosummary/mlx.core.reshape.html
index 7fd4d22f5..8d83b8a00 100644
--- a/docs/build/html/python/_autosummary/mlx.core.reshape.html
+++ b/docs/build/html/python/_autosummary/mlx.core.reshape.html
@@ -9,7 +9,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>mlx.core.reshape &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.reshape &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.rsqrt.html b/docs/build/html/python/_autosummary/mlx.core.rsqrt.html
index 5624cdba7..7559270a2 100644
--- a/docs/build/html/python/_autosummary/mlx.core.rsqrt.html
+++ b/docs/build/html/python/_autosummary/mlx.core.rsqrt.html
@@ -9,7 +9,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>mlx.core.rsqrt &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.rsqrt &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.save.html b/docs/build/html/python/_autosummary/mlx.core.save.html
index e88355724..0daf6534f 100644
--- a/docs/build/html/python/_autosummary/mlx.core.save.html
+++ b/docs/build/html/python/_autosummary/mlx.core.save.html
@@ -9,7 +9,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>mlx.core.save &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.save &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.savez.html b/docs/build/html/python/_autosummary/mlx.core.savez.html
index 054416432..7482b487e 100644
--- a/docs/build/html/python/_autosummary/mlx.core.savez.html
+++ b/docs/build/html/python/_autosummary/mlx.core.savez.html
@@ -9,7 +9,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>mlx.core.savez &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.savez &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.savez_compressed.html b/docs/build/html/python/_autosummary/mlx.core.savez_compressed.html
index 667654e99..ada2b82f3 100644
--- a/docs/build/html/python/_autosummary/mlx.core.savez_compressed.html
+++ b/docs/build/html/python/_autosummary/mlx.core.savez_compressed.html
@@ -9,7 +9,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>mlx.core.savez_compressed &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.savez_compressed &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.set_default_device.html b/docs/build/html/python/_autosummary/mlx.core.set_default_device.html
index 5bdfe5d64..472e35821 100644
--- a/docs/build/html/python/_autosummary/mlx.core.set_default_device.html
+++ b/docs/build/html/python/_autosummary/mlx.core.set_default_device.html
@@ -9,7 +9,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>mlx.core.set_default_device &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.set_default_device &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.set_default_stream.html b/docs/build/html/python/_autosummary/mlx.core.set_default_stream.html
index b33c6c844..cf86dd35d 100644
--- a/docs/build/html/python/_autosummary/mlx.core.set_default_stream.html
+++ b/docs/build/html/python/_autosummary/mlx.core.set_default_stream.html
@@ -9,7 +9,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>mlx.core.set_default_stream &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.set_default_stream &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.sigmoid.html b/docs/build/html/python/_autosummary/mlx.core.sigmoid.html
index 137a7014f..b87654a7e 100644
--- a/docs/build/html/python/_autosummary/mlx.core.sigmoid.html
+++ b/docs/build/html/python/_autosummary/mlx.core.sigmoid.html
@@ -9,7 +9,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>mlx.core.sigmoid &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.sigmoid &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -135,8 +135,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.sign.html b/docs/build/html/python/_autosummary/mlx.core.sign.html
index 8da467172..122a239d2 100644
--- a/docs/build/html/python/_autosummary/mlx.core.sign.html
+++ b/docs/build/html/python/_autosummary/mlx.core.sign.html
@@ -9,7 +9,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>mlx.core.sign &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.sign &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.sin.html b/docs/build/html/python/_autosummary/mlx.core.sin.html
index cd8a5da3f..b8586ad7d 100644
--- a/docs/build/html/python/_autosummary/mlx.core.sin.html
+++ b/docs/build/html/python/_autosummary/mlx.core.sin.html
@@ -9,7 +9,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>mlx.core.sin &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.sin &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.sinh.html b/docs/build/html/python/_autosummary/mlx.core.sinh.html
index 717bf9a5b..13c0273c3 100644
--- a/docs/build/html/python/_autosummary/mlx.core.sinh.html
+++ b/docs/build/html/python/_autosummary/mlx.core.sinh.html
@@ -9,7 +9,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>mlx.core.sinh &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.sinh &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.softmax.html b/docs/build/html/python/_autosummary/mlx.core.softmax.html
index 1716d8f75..7cda89ac7 100644
--- a/docs/build/html/python/_autosummary/mlx.core.softmax.html
+++ b/docs/build/html/python/_autosummary/mlx.core.softmax.html
@@ -9,7 +9,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>mlx.core.softmax &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.softmax &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.sort.html b/docs/build/html/python/_autosummary/mlx.core.sort.html
index a7279c31e..6be46f795 100644
--- a/docs/build/html/python/_autosummary/mlx.core.sort.html
+++ b/docs/build/html/python/_autosummary/mlx.core.sort.html
@@ -9,7 +9,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>mlx.core.sort &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.sort &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.split.html b/docs/build/html/python/_autosummary/mlx.core.split.html
index ea25bb9e6..c8c764f78 100644
--- a/docs/build/html/python/_autosummary/mlx.core.split.html
+++ b/docs/build/html/python/_autosummary/mlx.core.split.html
@@ -9,7 +9,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>mlx.core.split &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.split &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.sqrt.html b/docs/build/html/python/_autosummary/mlx.core.sqrt.html
index cfd7a21f1..18ac52427 100644
--- a/docs/build/html/python/_autosummary/mlx.core.sqrt.html
+++ b/docs/build/html/python/_autosummary/mlx.core.sqrt.html
@@ -9,7 +9,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>mlx.core.sqrt &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.sqrt &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.square.html b/docs/build/html/python/_autosummary/mlx.core.square.html
index ff718780c..c384891ff 100644
--- a/docs/build/html/python/_autosummary/mlx.core.square.html
+++ b/docs/build/html/python/_autosummary/mlx.core.square.html
@@ -9,7 +9,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>mlx.core.square &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.square &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.squeeze.html b/docs/build/html/python/_autosummary/mlx.core.squeeze.html
index 316723188..34aa38fda 100644
--- a/docs/build/html/python/_autosummary/mlx.core.squeeze.html
+++ b/docs/build/html/python/_autosummary/mlx.core.squeeze.html
@@ -9,7 +9,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>mlx.core.squeeze &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.squeeze &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.stop_gradient.html b/docs/build/html/python/_autosummary/mlx.core.stop_gradient.html
index 91e496f78..9a646af96 100644
--- a/docs/build/html/python/_autosummary/mlx.core.stop_gradient.html
+++ b/docs/build/html/python/_autosummary/mlx.core.stop_gradient.html
@@ -9,7 +9,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>mlx.core.stop_gradient &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.stop_gradient &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.subtract.html b/docs/build/html/python/_autosummary/mlx.core.subtract.html
index e7c53932c..901467252 100644
--- a/docs/build/html/python/_autosummary/mlx.core.subtract.html
+++ b/docs/build/html/python/_autosummary/mlx.core.subtract.html
@@ -9,7 +9,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>mlx.core.subtract &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.subtract &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.sum.html b/docs/build/html/python/_autosummary/mlx.core.sum.html
index bc155af3a..baa96c4f1 100644
--- a/docs/build/html/python/_autosummary/mlx.core.sum.html
+++ b/docs/build/html/python/_autosummary/mlx.core.sum.html
@@ -9,7 +9,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>mlx.core.sum &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.sum &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.take.html b/docs/build/html/python/_autosummary/mlx.core.take.html
index 005484c63..b9c1081a8 100644
--- a/docs/build/html/python/_autosummary/mlx.core.take.html
+++ b/docs/build/html/python/_autosummary/mlx.core.take.html
@@ -9,7 +9,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>mlx.core.take &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.take &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.take_along_axis.html b/docs/build/html/python/_autosummary/mlx.core.take_along_axis.html
index 309b69e6f..b1c26de59 100644
--- a/docs/build/html/python/_autosummary/mlx.core.take_along_axis.html
+++ b/docs/build/html/python/_autosummary/mlx.core.take_along_axis.html
@@ -9,7 +9,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>mlx.core.take_along_axis &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.take_along_axis &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.tan.html b/docs/build/html/python/_autosummary/mlx.core.tan.html
index b03cfc105..242a88f26 100644
--- a/docs/build/html/python/_autosummary/mlx.core.tan.html
+++ b/docs/build/html/python/_autosummary/mlx.core.tan.html
@@ -9,7 +9,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>mlx.core.tan &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.tan &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.tanh.html b/docs/build/html/python/_autosummary/mlx.core.tanh.html
index 025e073a7..fd82da57d 100644
--- a/docs/build/html/python/_autosummary/mlx.core.tanh.html
+++ b/docs/build/html/python/_autosummary/mlx.core.tanh.html
@@ -9,7 +9,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>mlx.core.tanh &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.tanh &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.transpose.html b/docs/build/html/python/_autosummary/mlx.core.transpose.html
index b123ef0c1..da43fad72 100644
--- a/docs/build/html/python/_autosummary/mlx.core.transpose.html
+++ b/docs/build/html/python/_autosummary/mlx.core.transpose.html
@@ -9,7 +9,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>mlx.core.transpose &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.transpose &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.value_and_grad.html b/docs/build/html/python/_autosummary/mlx.core.value_and_grad.html
index 9e0edcd8d..6b127894f 100644
--- a/docs/build/html/python/_autosummary/mlx.core.value_and_grad.html
+++ b/docs/build/html/python/_autosummary/mlx.core.value_and_grad.html
@@ -9,7 +9,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>mlx.core.value_and_grad &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.value_and_grad &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
@@ -542,7 +542,7 @@ document.write(`
 <h1>mlx.core.value_and_grad<a class="headerlink" href="#mlx-core-value-and-grad" title="Permalink to this heading">#</a></h1>
 <dl class="py function">
 <dt class="sig sig-object py" id="mlx.core.value_and_grad">
-<span class="sig-prename descclassname"><span class="pre">mlx.core.</span></span><span class="sig-name descname"><span class="pre">value_and_grad</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">fun</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">function</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">argnums</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">Union</span><span class="p"><span class="pre">[</span></span><a class="reference external" href="https://docs.python.org/3/library/functions.html#int" title="(in Python v3.12)"><span class="pre">int</span></a><span class="p"><span class="pre">,</span></span><span class="w"> </span><span class="pre">List</span><span class="p"><span class="pre">[</span></span><a class="reference external" href="https://docs.python.org/3/library/functions.html#int" title="(in Python v3.12)"><span class="pre">int</span></a><span class="p"><span class="pre">]</span></span><span class="p"><span class="pre">]</span></span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">0</span></span></em><span class="sig-paren">)</span> <span class="sig-return"><span class="sig-return-icon">&#x2192;</span> <span class="sig-return-typehint"><span class="pre">cpp_function</span></span></span><a class="headerlink" href="#mlx.core.value_and_grad" title="Permalink to this definition">#</a></dt>
+<span class="sig-prename descclassname"><span class="pre">mlx.core.</span></span><span class="sig-name descname"><span class="pre">value_and_grad</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">fun</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">function</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">argnums</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">Optional</span><span class="p"><span class="pre">[</span></span><span class="pre">Union</span><span class="p"><span class="pre">[</span></span><a class="reference external" href="https://docs.python.org/3/library/functions.html#int" title="(in Python v3.12)"><span class="pre">int</span></a><span class="p"><span class="pre">,</span></span><span class="w"> </span><span class="pre">List</span><span class="p"><span class="pre">[</span></span><a class="reference external" href="https://docs.python.org/3/library/functions.html#int" title="(in Python v3.12)"><span class="pre">int</span></a><span class="p"><span class="pre">]</span></span><span class="p"><span class="pre">]</span></span><span class="p"><span class="pre">]</span></span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">None</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">argnames</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">Union</span><span class="p"><span class="pre">[</span></span><a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.12)"><span class="pre">str</span></a><span class="p"><span class="pre">,</span></span><span class="w"> </span><span class="pre">List</span><span class="p"><span class="pre">[</span></span><a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.12)"><span class="pre">str</span></a><span class="p"><span class="pre">]</span></span><span class="p"><span class="pre">]</span></span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">[]</span></span></em><span class="sig-paren">)</span> <span class="sig-return"><span class="sig-return-icon">&#x2192;</span> <span class="sig-return-typehint"><span class="pre">cpp_function</span></span></span><a class="headerlink" href="#mlx.core.value_and_grad" title="Permalink to this definition">#</a></dt>
 <dd><p>Returns a function which computes the value and gradient of <code class="docutils literal notranslate"><span class="pre">fun</span></code>.</p>
 <p>The function passed to <a class="reference internal" href="#mlx.core.value_and_grad" title="mlx.core.value_and_grad"><code class="xref py py-func docutils literal notranslate"><span class="pre">value_and_grad()</span></code></a> should return either
 a scalar loss or a tuple in which the first element is a scalar
@@ -577,8 +577,13 @@ loss and the remaining elements can be anything.</p>
 a scalar output <a class="reference internal" href="mlx.core.array.html#mlx.core.array" title="mlx.core.array"><code class="xref py py-class docutils literal notranslate"><span class="pre">array</span></code></a> or a tuple the first element
 of which should be a scalar <a class="reference internal" href="mlx.core.array.html#mlx.core.array" title="mlx.core.array"><code class="xref py py-class docutils literal notranslate"><span class="pre">array</span></code></a>.</p></li>
 <li><p><strong>argnums</strong> (<a class="reference external" href="https://docs.python.org/3/library/functions.html#int" title="(in Python v3.12)"><em>int</em></a><em> or </em><a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#list" title="(in Python v3.12)"><em>list</em></a><em>(</em><a class="reference external" href="https://docs.python.org/3/library/functions.html#int" title="(in Python v3.12)"><em>int</em></a><em>)</em><em>, </em><em>optional</em>) – Specify the index (or indices)
-of the arguments of <code class="docutils literal notranslate"><span class="pre">fun</span></code> to compute the gradient with respect to. If
-unspecified this defaults to <code class="docutils literal notranslate"><span class="pre">0</span></code> indicating <code class="docutils literal notranslate"><span class="pre">fun</span></code>’s first argument.</p></li>
+of the positional arguments of <code class="docutils literal notranslate"><span class="pre">fun</span></code> to compute the gradient
+with respect to. If neither <code class="docutils literal notranslate"><span class="pre">argnums</span></code> nor <code class="docutils literal notranslate"><span class="pre">argnames</span></code> are
+provided <code class="docutils literal notranslate"><span class="pre">argnums</span></code> defaults to <code class="docutils literal notranslate"><span class="pre">0</span></code> indicating <code class="docutils literal notranslate"><span class="pre">fun</span></code>’s first
+argument.</p></li>
+<li><p><strong>argnames</strong> (<a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.12)"><em>str</em></a><em> or </em><a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#list" title="(in Python v3.12)"><em>list</em></a><em>(</em><a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.12)"><em>str</em></a><em>)</em><em>, </em><em>optional</em>) – Specify keyword arguments of
+<code class="docutils literal notranslate"><span class="pre">fun</span></code> to compute gradients with respect to. It defaults to [] so
+no gradients for keyword arguments by default.</p></li>
 </ul>
 </dd>
 <dt class="field-even">Returns<span class="colon">:</span></dt>
diff --git a/docs/build/html/python/_autosummary/mlx.core.var.html b/docs/build/html/python/_autosummary/mlx.core.var.html
index 0cb7a0b60..288524e35 100644
--- a/docs/build/html/python/_autosummary/mlx.core.var.html
+++ b/docs/build/html/python/_autosummary/mlx.core.var.html
@@ -9,7 +9,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>mlx.core.var &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.var &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.vjp.html b/docs/build/html/python/_autosummary/mlx.core.vjp.html
index 25fc193aa..29627cf73 100644
--- a/docs/build/html/python/_autosummary/mlx.core.vjp.html
+++ b/docs/build/html/python/_autosummary/mlx.core.vjp.html
@@ -9,7 +9,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>mlx.core.vjp &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.vjp &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.vmap.html b/docs/build/html/python/_autosummary/mlx.core.vmap.html
index 7a214f619..b471d3747 100644
--- a/docs/build/html/python/_autosummary/mlx.core.vmap.html
+++ b/docs/build/html/python/_autosummary/mlx.core.vmap.html
@@ -9,7 +9,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>mlx.core.vmap &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.vmap &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.where.html b/docs/build/html/python/_autosummary/mlx.core.where.html
index cab721c8e..cccf59fce 100644
--- a/docs/build/html/python/_autosummary/mlx.core.where.html
+++ b/docs/build/html/python/_autosummary/mlx.core.where.html
@@ -9,7 +9,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>mlx.core.where &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.where &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.zeros.html b/docs/build/html/python/_autosummary/mlx.core.zeros.html
index 95733803f..2e46addce 100644
--- a/docs/build/html/python/_autosummary/mlx.core.zeros.html
+++ b/docs/build/html/python/_autosummary/mlx.core.zeros.html
@@ -9,7 +9,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>mlx.core.zeros &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.zeros &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.core.zeros_like.html b/docs/build/html/python/_autosummary/mlx.core.zeros_like.html
index a5a214c48..7a57c8e0c 100644
--- a/docs/build/html/python/_autosummary/mlx.core.zeros_like.html
+++ b/docs/build/html/python/_autosummary/mlx.core.zeros_like.html
@@ -9,7 +9,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>mlx.core.zeros_like &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.core.zeros_like &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.nn.Conv1d.html b/docs/build/html/python/_autosummary/mlx.nn.Conv1d.html
index f680d9f4c..80e4f08a7 100644
--- a/docs/build/html/python/_autosummary/mlx.nn.Conv1d.html
+++ b/docs/build/html/python/_autosummary/mlx.nn.Conv1d.html
@@ -9,7 +9,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>mlx.nn.Conv1d &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.nn.Conv1d &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.nn.Conv2d.html b/docs/build/html/python/_autosummary/mlx.nn.Conv2d.html
index 820962bca..01d617ca7 100644
--- a/docs/build/html/python/_autosummary/mlx.nn.Conv2d.html
+++ b/docs/build/html/python/_autosummary/mlx.nn.Conv2d.html
@@ -9,7 +9,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>mlx.nn.Conv2d &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.nn.Conv2d &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.nn.Embedding.html b/docs/build/html/python/_autosummary/mlx.nn.Embedding.html
index 010bcd4b8..f2f01f4a8 100644
--- a/docs/build/html/python/_autosummary/mlx.nn.Embedding.html
+++ b/docs/build/html/python/_autosummary/mlx.nn.Embedding.html
@@ -9,7 +9,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>mlx.nn.Embedding &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.nn.Embedding &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.nn.GELU.html b/docs/build/html/python/_autosummary/mlx.nn.GELU.html
index d662a3823..629322841 100644
--- a/docs/build/html/python/_autosummary/mlx.nn.GELU.html
+++ b/docs/build/html/python/_autosummary/mlx.nn.GELU.html
@@ -9,7 +9,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>mlx.nn.GELU &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.nn.GELU &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -135,8 +135,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.nn.GroupNorm.html b/docs/build/html/python/_autosummary/mlx.nn.GroupNorm.html
index 4c33567f9..591fbb581 100644
--- a/docs/build/html/python/_autosummary/mlx.nn.GroupNorm.html
+++ b/docs/build/html/python/_autosummary/mlx.nn.GroupNorm.html
@@ -9,7 +9,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>mlx.nn.GroupNorm &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.nn.GroupNorm &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -135,8 +135,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.nn.LayerNorm.html b/docs/build/html/python/_autosummary/mlx.nn.LayerNorm.html
index ab6153574..bca6e9220 100644
--- a/docs/build/html/python/_autosummary/mlx.nn.LayerNorm.html
+++ b/docs/build/html/python/_autosummary/mlx.nn.LayerNorm.html
@@ -9,7 +9,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>mlx.nn.LayerNorm &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.nn.LayerNorm &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -135,8 +135,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.nn.Linear.html b/docs/build/html/python/_autosummary/mlx.nn.Linear.html
index c52c1c410..4ab2cad00 100644
--- a/docs/build/html/python/_autosummary/mlx.nn.Linear.html
+++ b/docs/build/html/python/_autosummary/mlx.nn.Linear.html
@@ -9,7 +9,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>mlx.nn.Linear &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.nn.Linear &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.nn.MultiHeadAttention.html b/docs/build/html/python/_autosummary/mlx.nn.MultiHeadAttention.html
index 14e58ab97..44b63ca55 100644
--- a/docs/build/html/python/_autosummary/mlx.nn.MultiHeadAttention.html
+++ b/docs/build/html/python/_autosummary/mlx.nn.MultiHeadAttention.html
@@ -9,7 +9,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>mlx.nn.MultiHeadAttention &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.nn.MultiHeadAttention &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.nn.RMSNorm.html b/docs/build/html/python/_autosummary/mlx.nn.RMSNorm.html
index d5dec4a41..252a6ac61 100644
--- a/docs/build/html/python/_autosummary/mlx.nn.RMSNorm.html
+++ b/docs/build/html/python/_autosummary/mlx.nn.RMSNorm.html
@@ -9,7 +9,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>mlx.nn.RMSNorm &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.nn.RMSNorm &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -135,8 +135,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.nn.ReLU.html b/docs/build/html/python/_autosummary/mlx.nn.ReLU.html
index 47dbff859..150402c17 100644
--- a/docs/build/html/python/_autosummary/mlx.nn.ReLU.html
+++ b/docs/build/html/python/_autosummary/mlx.nn.ReLU.html
@@ -9,7 +9,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>mlx.nn.ReLU &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.nn.ReLU &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.nn.RoPE.html b/docs/build/html/python/_autosummary/mlx.nn.RoPE.html
index bd65e6cb5..a599f0993 100644
--- a/docs/build/html/python/_autosummary/mlx.nn.RoPE.html
+++ b/docs/build/html/python/_autosummary/mlx.nn.RoPE.html
@@ -9,7 +9,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>mlx.nn.RoPE &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.nn.RoPE &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.nn.Sequential.html b/docs/build/html/python/_autosummary/mlx.nn.Sequential.html
index 0807c718e..c52d319b7 100644
--- a/docs/build/html/python/_autosummary/mlx.nn.Sequential.html
+++ b/docs/build/html/python/_autosummary/mlx.nn.Sequential.html
@@ -9,7 +9,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>mlx.nn.Sequential &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.nn.Sequential &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.nn.SiLU.html b/docs/build/html/python/_autosummary/mlx.nn.SiLU.html
index f391f12d1..5f96bf11b 100644
--- a/docs/build/html/python/_autosummary/mlx.nn.SiLU.html
+++ b/docs/build/html/python/_autosummary/mlx.nn.SiLU.html
@@ -9,7 +9,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>mlx.nn.SiLU &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.nn.SiLU &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -135,8 +135,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.nn.value_and_grad.html b/docs/build/html/python/_autosummary/mlx.nn.value_and_grad.html
index 9a6f401fb..b6855afd5 100644
--- a/docs/build/html/python/_autosummary/mlx.nn.value_and_grad.html
+++ b/docs/build/html/python/_autosummary/mlx.nn.value_and_grad.html
@@ -9,7 +9,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>mlx.nn.value_and_grad &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.nn.value_and_grad &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.optimizers.Adam.html b/docs/build/html/python/_autosummary/mlx.optimizers.Adam.html
index 4a6cbf8b7..e9c1bff46 100644
--- a/docs/build/html/python/_autosummary/mlx.optimizers.Adam.html
+++ b/docs/build/html/python/_autosummary/mlx.optimizers.Adam.html
@@ -9,7 +9,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>mlx.optimizers.Adam &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.optimizers.Adam &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -135,8 +135,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.optimizers.Optimizer.html b/docs/build/html/python/_autosummary/mlx.optimizers.Optimizer.html
index 24ca0a80b..81a66ab17 100644
--- a/docs/build/html/python/_autosummary/mlx.optimizers.Optimizer.html
+++ b/docs/build/html/python/_autosummary/mlx.optimizers.Optimizer.html
@@ -9,7 +9,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>mlx.optimizers.Optimizer &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.optimizers.Optimizer &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.optimizers.OptimizerState.html b/docs/build/html/python/_autosummary/mlx.optimizers.OptimizerState.html
index 51589e9f5..892205e0d 100644
--- a/docs/build/html/python/_autosummary/mlx.optimizers.OptimizerState.html
+++ b/docs/build/html/python/_autosummary/mlx.optimizers.OptimizerState.html
@@ -9,7 +9,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>mlx.optimizers.OptimizerState &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.optimizers.OptimizerState &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.optimizers.SGD.html b/docs/build/html/python/_autosummary/mlx.optimizers.SGD.html
index fd4468408..676e6e75e 100644
--- a/docs/build/html/python/_autosummary/mlx.optimizers.SGD.html
+++ b/docs/build/html/python/_autosummary/mlx.optimizers.SGD.html
@@ -9,7 +9,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>mlx.optimizers.SGD &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.optimizers.SGD &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -135,8 +135,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.utils.tree_flatten.html b/docs/build/html/python/_autosummary/mlx.utils.tree_flatten.html
index 7f20e5b16..fcbba82a3 100644
--- a/docs/build/html/python/_autosummary/mlx.utils.tree_flatten.html
+++ b/docs/build/html/python/_autosummary/mlx.utils.tree_flatten.html
@@ -9,7 +9,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>mlx.utils.tree_flatten &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.utils.tree_flatten &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.utils.tree_map.html b/docs/build/html/python/_autosummary/mlx.utils.tree_map.html
index a0cfc632d..70916c884 100644
--- a/docs/build/html/python/_autosummary/mlx.utils.tree_map.html
+++ b/docs/build/html/python/_autosummary/mlx.utils.tree_map.html
@@ -9,7 +9,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>mlx.utils.tree_map &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.utils.tree_map &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary/mlx.utils.tree_unflatten.html b/docs/build/html/python/_autosummary/mlx.utils.tree_unflatten.html
index 2a2a8b769..fc5debda3 100644
--- a/docs/build/html/python/_autosummary/mlx.utils.tree_unflatten.html
+++ b/docs/build/html/python/_autosummary/mlx.utils.tree_unflatten.html
@@ -9,7 +9,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>mlx.utils.tree_unflatten &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.utils.tree_unflatten &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary_functions/mlx.nn.gelu.html b/docs/build/html/python/_autosummary_functions/mlx.nn.gelu.html
index 456b85cd0..bdb6ed314 100644
--- a/docs/build/html/python/_autosummary_functions/mlx.nn.gelu.html
+++ b/docs/build/html/python/_autosummary_functions/mlx.nn.gelu.html
@@ -9,7 +9,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>mlx.nn.gelu &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.nn.gelu &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -135,8 +135,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary_functions/mlx.nn.gelu_approx.html b/docs/build/html/python/_autosummary_functions/mlx.nn.gelu_approx.html
index 1b308acfe..786ad1961 100644
--- a/docs/build/html/python/_autosummary_functions/mlx.nn.gelu_approx.html
+++ b/docs/build/html/python/_autosummary_functions/mlx.nn.gelu_approx.html
@@ -9,7 +9,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>mlx.nn.gelu_approx &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.nn.gelu_approx &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -135,8 +135,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary_functions/mlx.nn.gelu_fast_approx.html b/docs/build/html/python/_autosummary_functions/mlx.nn.gelu_fast_approx.html
index 386588e82..482ac40e9 100644
--- a/docs/build/html/python/_autosummary_functions/mlx.nn.gelu_fast_approx.html
+++ b/docs/build/html/python/_autosummary_functions/mlx.nn.gelu_fast_approx.html
@@ -9,7 +9,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>mlx.nn.gelu_fast_approx &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.nn.gelu_fast_approx &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -135,8 +135,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary_functions/mlx.nn.relu.html b/docs/build/html/python/_autosummary_functions/mlx.nn.relu.html
index 828b99035..c85d2922b 100644
--- a/docs/build/html/python/_autosummary_functions/mlx.nn.relu.html
+++ b/docs/build/html/python/_autosummary_functions/mlx.nn.relu.html
@@ -9,7 +9,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>mlx.nn.relu &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.nn.relu &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/_autosummary_functions/mlx.nn.silu.html b/docs/build/html/python/_autosummary_functions/mlx.nn.silu.html
index 1a46bb6ae..cff4a5e6b 100644
--- a/docs/build/html/python/_autosummary_functions/mlx.nn.silu.html
+++ b/docs/build/html/python/_autosummary_functions/mlx.nn.silu.html
@@ -9,7 +9,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>mlx.nn.silu &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.nn.silu &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -135,8 +135,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/array.html b/docs/build/html/python/array.html
index 8997ff49b..3ef6b08f9 100644
--- a/docs/build/html/python/array.html
+++ b/docs/build/html/python/array.html
@@ -9,7 +9,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>Array &#8212; MLX 0.0.3 documentation</title>
+    <title>Array &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/data_types.html b/docs/build/html/python/data_types.html
index a6d545a6e..f259dbfab 100644
--- a/docs/build/html/python/data_types.html
+++ b/docs/build/html/python/data_types.html
@@ -9,7 +9,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>Data Types &#8212; MLX 0.0.3 documentation</title>
+    <title>Data Types &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -132,8 +132,8 @@
       
     
     
-    <img src="../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/devices_and_streams.html b/docs/build/html/python/devices_and_streams.html
index 2d2c5db13..2f41816a1 100644
--- a/docs/build/html/python/devices_and_streams.html
+++ b/docs/build/html/python/devices_and_streams.html
@@ -9,7 +9,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>Devices and Streams &#8212; MLX 0.0.3 documentation</title>
+    <title>Devices and Streams &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/fft.html b/docs/build/html/python/fft.html
index 766385dbd..f25dadc64 100644
--- a/docs/build/html/python/fft.html
+++ b/docs/build/html/python/fft.html
@@ -9,7 +9,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>FFT &#8212; MLX 0.0.3 documentation</title>
+    <title>FFT &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/nn.html b/docs/build/html/python/nn.html
index 13914f3cc..9e49a3113 100644
--- a/docs/build/html/python/nn.html
+++ b/docs/build/html/python/nn.html
@@ -9,7 +9,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>Neural Networks &#8212; MLX 0.0.3 documentation</title>
+    <title>Neural Networks &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/nn/module.html b/docs/build/html/python/nn/module.html
index 0d916230d..b4f5c2ac3 100644
--- a/docs/build/html/python/nn/module.html
+++ b/docs/build/html/python/nn/module.html
@@ -9,7 +9,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>mlx.nn.Module &#8212; MLX 0.0.3 documentation</title>
+    <title>mlx.nn.Module &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -132,8 +132,8 @@
       
     
     
-    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/ops.html b/docs/build/html/python/ops.html
index 7254f765e..9f301399b 100644
--- a/docs/build/html/python/ops.html
+++ b/docs/build/html/python/ops.html
@@ -9,7 +9,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>Operations &#8212; MLX 0.0.3 documentation</title>
+    <title>Operations &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/optimizers.html b/docs/build/html/python/optimizers.html
index 9f8070ac9..8c0edb26c 100644
--- a/docs/build/html/python/optimizers.html
+++ b/docs/build/html/python/optimizers.html
@@ -9,7 +9,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.0.3 documentation</title>
+    <title>Optimizers &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/random.html b/docs/build/html/python/random.html
index ad5fae1e4..b6ecb6025 100644
--- a/docs/build/html/python/random.html
+++ b/docs/build/html/python/random.html
@@ -9,7 +9,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>Random &#8212; MLX 0.0.3 documentation</title>
+    <title>Random &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/python/transforms.html b/docs/build/html/python/transforms.html
index e74926bc9..3ba1989a8 100644
--- a/docs/build/html/python/transforms.html
+++ b/docs/build/html/python/transforms.html
@@ -9,7 +9,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>Transforms &#8212; MLX 0.0.3 documentation</title>
+    <title>Transforms &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
@@ -545,10 +545,10 @@ document.write(`
 <tr class="row-odd"><td><p><a class="reference internal" href="_autosummary/mlx.core.eval.html#mlx.core.eval" title="mlx.core.eval"><code class="xref py py-obj docutils literal notranslate"><span class="pre">eval</span></code></a>(*args[, retain_graph])</p></td>
 <td><p>Evaluate an <a class="reference internal" href="_autosummary/mlx.core.array.html#mlx.core.array" title="mlx.core.array"><code class="xref py py-class docutils literal notranslate"><span class="pre">array</span></code></a> or tree of <a class="reference internal" href="_autosummary/mlx.core.array.html#mlx.core.array" title="mlx.core.array"><code class="xref py py-class docutils literal notranslate"><span class="pre">array</span></code></a>.</p></td>
 </tr>
-<tr class="row-even"><td><p><a class="reference internal" href="_autosummary/mlx.core.grad.html#mlx.core.grad" title="mlx.core.grad"><code class="xref py py-obj docutils literal notranslate"><span class="pre">grad</span></code></a>(fun[, argnums])</p></td>
+<tr class="row-even"><td><p><a class="reference internal" href="_autosummary/mlx.core.grad.html#mlx.core.grad" title="mlx.core.grad"><code class="xref py py-obj docutils literal notranslate"><span class="pre">grad</span></code></a>(fun[, argnums, argnames])</p></td>
 <td><p>Returns a function which computes the gradient of <code class="docutils literal notranslate"><span class="pre">fun</span></code>.</p></td>
 </tr>
-<tr class="row-odd"><td><p><a class="reference internal" href="_autosummary/mlx.core.value_and_grad.html#mlx.core.value_and_grad" title="mlx.core.value_and_grad"><code class="xref py py-obj docutils literal notranslate"><span class="pre">value_and_grad</span></code></a>(fun[, argnums])</p></td>
+<tr class="row-odd"><td><p><a class="reference internal" href="_autosummary/mlx.core.value_and_grad.html#mlx.core.value_and_grad" title="mlx.core.value_and_grad"><code class="xref py py-obj docutils literal notranslate"><span class="pre">value_and_grad</span></code></a>(fun[, argnums, argnames])</p></td>
 <td><p>Returns a function which computes the value and gradient of <code class="docutils literal notranslate"><span class="pre">fun</span></code>.</p></td>
 </tr>
 <tr class="row-even"><td><p><a class="reference internal" href="_autosummary/mlx.core.jvp.html#mlx.core.jvp" title="mlx.core.jvp"><code class="xref py py-obj docutils literal notranslate"><span class="pre">jvp</span></code></a>(fun, primals, tangents)</p></td>
diff --git a/docs/build/html/python/tree_utils.html b/docs/build/html/python/tree_utils.html
index 7442d90e5..2e108e049 100644
--- a/docs/build/html/python/tree_utils.html
+++ b/docs/build/html/python/tree_utils.html
@@ -9,7 +9,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>Tree Utils &#8212; MLX 0.0.3 documentation</title>
+    <title>Tree Utils &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="../_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/quick_start.html b/docs/build/html/quick_start.html
index b1e85805f..12ba9c909 100644
--- a/docs/build/html/quick_start.html
+++ b/docs/build/html/quick_start.html
@@ -9,7 +9,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>Quick Start Guide &#8212; MLX 0.0.3 documentation</title>
+    <title>Quick Start Guide &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/search.html b/docs/build/html/search.html
index 52b2867b2..a817da5d1 100644
--- a/docs/build/html/search.html
+++ b/docs/build/html/search.html
@@ -7,7 +7,7 @@
 
   <head>
     <meta charset="utf-8" />
-    <meta name="viewport" content="width=device-width, initial-scale=1.0" /><title>Search - MLX 0.0.3 documentation</title>
+    <meta name="viewport" content="width=device-width, initial-scale=1.0" /><title>Search - MLX 0.0.4 documentation</title>
   
   
   
@@ -133,8 +133,8 @@
       
     
     
-    <img src="_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>
diff --git a/docs/build/html/searchindex.js b/docs/build/html/searchindex.js
index 4e4dc94db..ac2ed5b34 100644
--- a/docs/build/html/searchindex.js
+++ b/docs/build/html/searchindex.js
@@ -1 +1 @@
-Search.setIndex({"docnames": ["cpp/ops", "dev/extensions", "examples/linear_regression", "examples/llama-inference", "examples/mlp", "index", "install", "python/_autosummary/mlx.core.Device", "python/_autosummary/mlx.core.Dtype", "python/_autosummary/mlx.core.Stream", "python/_autosummary/mlx.core.abs", "python/_autosummary/mlx.core.add", "python/_autosummary/mlx.core.all", "python/_autosummary/mlx.core.allclose", "python/_autosummary/mlx.core.any", "python/_autosummary/mlx.core.arange", "python/_autosummary/mlx.core.arccos", "python/_autosummary/mlx.core.arccosh", "python/_autosummary/mlx.core.arcsin", "python/_autosummary/mlx.core.arcsinh", "python/_autosummary/mlx.core.arctan", "python/_autosummary/mlx.core.arctanh", "python/_autosummary/mlx.core.argmax", "python/_autosummary/mlx.core.argmin", "python/_autosummary/mlx.core.argpartition", "python/_autosummary/mlx.core.argsort", "python/_autosummary/mlx.core.array", "python/_autosummary/mlx.core.array.T", "python/_autosummary/mlx.core.array.abs", "python/_autosummary/mlx.core.array.all", "python/_autosummary/mlx.core.array.any", "python/_autosummary/mlx.core.array.argmax", "python/_autosummary/mlx.core.array.argmin", "python/_autosummary/mlx.core.array.astype", "python/_autosummary/mlx.core.array.cos", "python/_autosummary/mlx.core.array.dtype", "python/_autosummary/mlx.core.array.exp", "python/_autosummary/mlx.core.array.item", "python/_autosummary/mlx.core.array.log", "python/_autosummary/mlx.core.array.log1p", "python/_autosummary/mlx.core.array.logsumexp", "python/_autosummary/mlx.core.array.max", "python/_autosummary/mlx.core.array.mean", "python/_autosummary/mlx.core.array.min", "python/_autosummary/mlx.core.array.ndim", "python/_autosummary/mlx.core.array.prod", "python/_autosummary/mlx.core.array.reciprocal", "python/_autosummary/mlx.core.array.reshape", "python/_autosummary/mlx.core.array.rsqrt", "python/_autosummary/mlx.core.array.shape", "python/_autosummary/mlx.core.array.sin", "python/_autosummary/mlx.core.array.size", "python/_autosummary/mlx.core.array.split", "python/_autosummary/mlx.core.array.sqrt", "python/_autosummary/mlx.core.array.square", "python/_autosummary/mlx.core.array.sum", "python/_autosummary/mlx.core.array.tolist", "python/_autosummary/mlx.core.array.transpose", "python/_autosummary/mlx.core.array.var", "python/_autosummary/mlx.core.array_equal", "python/_autosummary/mlx.core.broadcast_to", "python/_autosummary/mlx.core.concatenate", "python/_autosummary/mlx.core.conv1d", "python/_autosummary/mlx.core.conv2d", "python/_autosummary/mlx.core.convolve", "python/_autosummary/mlx.core.cos", "python/_autosummary/mlx.core.cosh", "python/_autosummary/mlx.core.default_device", "python/_autosummary/mlx.core.default_stream", "python/_autosummary/mlx.core.divide", "python/_autosummary/mlx.core.equal", "python/_autosummary/mlx.core.erf", "python/_autosummary/mlx.core.erfinv", "python/_autosummary/mlx.core.eval", "python/_autosummary/mlx.core.exp", "python/_autosummary/mlx.core.expand_dims", "python/_autosummary/mlx.core.fft.fft", "python/_autosummary/mlx.core.fft.fft2", "python/_autosummary/mlx.core.fft.fftn", "python/_autosummary/mlx.core.fft.ifft", "python/_autosummary/mlx.core.fft.ifft2", "python/_autosummary/mlx.core.fft.ifftn", "python/_autosummary/mlx.core.fft.irfft", "python/_autosummary/mlx.core.fft.irfft2", "python/_autosummary/mlx.core.fft.irfftn", "python/_autosummary/mlx.core.fft.rfft", "python/_autosummary/mlx.core.fft.rfft2", "python/_autosummary/mlx.core.fft.rfftn", "python/_autosummary/mlx.core.full", "python/_autosummary/mlx.core.grad", "python/_autosummary/mlx.core.greater", "python/_autosummary/mlx.core.greater_equal", "python/_autosummary/mlx.core.jvp", "python/_autosummary/mlx.core.less", "python/_autosummary/mlx.core.less_equal", "python/_autosummary/mlx.core.load", "python/_autosummary/mlx.core.log", "python/_autosummary/mlx.core.log10", "python/_autosummary/mlx.core.log1p", "python/_autosummary/mlx.core.log2", "python/_autosummary/mlx.core.logaddexp", "python/_autosummary/mlx.core.logical_not", "python/_autosummary/mlx.core.logsumexp", "python/_autosummary/mlx.core.matmul", "python/_autosummary/mlx.core.max", "python/_autosummary/mlx.core.maximum", "python/_autosummary/mlx.core.mean", "python/_autosummary/mlx.core.min", "python/_autosummary/mlx.core.minimum", "python/_autosummary/mlx.core.multiply", "python/_autosummary/mlx.core.negative", "python/_autosummary/mlx.core.new_stream", "python/_autosummary/mlx.core.ones", "python/_autosummary/mlx.core.ones_like", "python/_autosummary/mlx.core.pad", "python/_autosummary/mlx.core.partition", "python/_autosummary/mlx.core.prod", "python/_autosummary/mlx.core.random.bernoulli", "python/_autosummary/mlx.core.random.categorical", "python/_autosummary/mlx.core.random.gumbel", "python/_autosummary/mlx.core.random.key", "python/_autosummary/mlx.core.random.normal", "python/_autosummary/mlx.core.random.randint", "python/_autosummary/mlx.core.random.seed", "python/_autosummary/mlx.core.random.split", "python/_autosummary/mlx.core.random.truncated_normal", "python/_autosummary/mlx.core.random.uniform", "python/_autosummary/mlx.core.reciprocal", "python/_autosummary/mlx.core.reshape", "python/_autosummary/mlx.core.rsqrt", "python/_autosummary/mlx.core.save", "python/_autosummary/mlx.core.savez", "python/_autosummary/mlx.core.savez_compressed", "python/_autosummary/mlx.core.set_default_device", "python/_autosummary/mlx.core.set_default_stream", "python/_autosummary/mlx.core.sigmoid", "python/_autosummary/mlx.core.sign", "python/_autosummary/mlx.core.sin", "python/_autosummary/mlx.core.sinh", "python/_autosummary/mlx.core.softmax", "python/_autosummary/mlx.core.sort", "python/_autosummary/mlx.core.split", "python/_autosummary/mlx.core.sqrt", "python/_autosummary/mlx.core.square", "python/_autosummary/mlx.core.squeeze", "python/_autosummary/mlx.core.stop_gradient", "python/_autosummary/mlx.core.subtract", "python/_autosummary/mlx.core.sum", "python/_autosummary/mlx.core.take", "python/_autosummary/mlx.core.take_along_axis", "python/_autosummary/mlx.core.tan", "python/_autosummary/mlx.core.tanh", "python/_autosummary/mlx.core.transpose", "python/_autosummary/mlx.core.value_and_grad", "python/_autosummary/mlx.core.var", "python/_autosummary/mlx.core.vjp", "python/_autosummary/mlx.core.vmap", "python/_autosummary/mlx.core.where", "python/_autosummary/mlx.core.zeros", "python/_autosummary/mlx.core.zeros_like", "python/_autosummary/mlx.nn.Conv1d", "python/_autosummary/mlx.nn.Conv2d", "python/_autosummary/mlx.nn.Embedding", "python/_autosummary/mlx.nn.GELU", "python/_autosummary/mlx.nn.GroupNorm", "python/_autosummary/mlx.nn.LayerNorm", "python/_autosummary/mlx.nn.Linear", "python/_autosummary/mlx.nn.MultiHeadAttention", "python/_autosummary/mlx.nn.RMSNorm", "python/_autosummary/mlx.nn.ReLU", "python/_autosummary/mlx.nn.RoPE", "python/_autosummary/mlx.nn.Sequential", "python/_autosummary/mlx.nn.SiLU", "python/_autosummary/mlx.nn.value_and_grad", "python/_autosummary/mlx.optimizers.Adam", "python/_autosummary/mlx.optimizers.Optimizer", "python/_autosummary/mlx.optimizers.OptimizerState", "python/_autosummary/mlx.optimizers.SGD", "python/_autosummary/mlx.utils.tree_flatten", "python/_autosummary/mlx.utils.tree_map", "python/_autosummary/mlx.utils.tree_unflatten", "python/_autosummary_functions/mlx.nn.gelu", "python/_autosummary_functions/mlx.nn.gelu_approx", "python/_autosummary_functions/mlx.nn.gelu_fast_approx", "python/_autosummary_functions/mlx.nn.relu", "python/_autosummary_functions/mlx.nn.silu", "python/array", "python/data_types", "python/devices_and_streams", "python/fft", "python/nn", "python/nn/module", "python/ops", "python/optimizers", "python/random", "python/transforms", "python/tree_utils", "quick_start", "using_streams"], "filenames": ["cpp/ops.rst", "dev/extensions.rst", "examples/linear_regression.rst", "examples/llama-inference.rst", "examples/mlp.rst", "index.rst", "install.rst", "python/_autosummary/mlx.core.Device.rst", "python/_autosummary/mlx.core.Dtype.rst", "python/_autosummary/mlx.core.Stream.rst", "python/_autosummary/mlx.core.abs.rst", "python/_autosummary/mlx.core.add.rst", "python/_autosummary/mlx.core.all.rst", "python/_autosummary/mlx.core.allclose.rst", "python/_autosummary/mlx.core.any.rst", "python/_autosummary/mlx.core.arange.rst", "python/_autosummary/mlx.core.arccos.rst", "python/_autosummary/mlx.core.arccosh.rst", "python/_autosummary/mlx.core.arcsin.rst", "python/_autosummary/mlx.core.arcsinh.rst", "python/_autosummary/mlx.core.arctan.rst", "python/_autosummary/mlx.core.arctanh.rst", "python/_autosummary/mlx.core.argmax.rst", "python/_autosummary/mlx.core.argmin.rst", "python/_autosummary/mlx.core.argpartition.rst", "python/_autosummary/mlx.core.argsort.rst", "python/_autosummary/mlx.core.array.rst", "python/_autosummary/mlx.core.array.T.rst", "python/_autosummary/mlx.core.array.abs.rst", "python/_autosummary/mlx.core.array.all.rst", "python/_autosummary/mlx.core.array.any.rst", "python/_autosummary/mlx.core.array.argmax.rst", "python/_autosummary/mlx.core.array.argmin.rst", "python/_autosummary/mlx.core.array.astype.rst", "python/_autosummary/mlx.core.array.cos.rst", "python/_autosummary/mlx.core.array.dtype.rst", "python/_autosummary/mlx.core.array.exp.rst", "python/_autosummary/mlx.core.array.item.rst", "python/_autosummary/mlx.core.array.log.rst", "python/_autosummary/mlx.core.array.log1p.rst", "python/_autosummary/mlx.core.array.logsumexp.rst", "python/_autosummary/mlx.core.array.max.rst", "python/_autosummary/mlx.core.array.mean.rst", "python/_autosummary/mlx.core.array.min.rst", "python/_autosummary/mlx.core.array.ndim.rst", "python/_autosummary/mlx.core.array.prod.rst", "python/_autosummary/mlx.core.array.reciprocal.rst", "python/_autosummary/mlx.core.array.reshape.rst", "python/_autosummary/mlx.core.array.rsqrt.rst", "python/_autosummary/mlx.core.array.shape.rst", "python/_autosummary/mlx.core.array.sin.rst", "python/_autosummary/mlx.core.array.size.rst", "python/_autosummary/mlx.core.array.split.rst", "python/_autosummary/mlx.core.array.sqrt.rst", "python/_autosummary/mlx.core.array.square.rst", "python/_autosummary/mlx.core.array.sum.rst", "python/_autosummary/mlx.core.array.tolist.rst", "python/_autosummary/mlx.core.array.transpose.rst", "python/_autosummary/mlx.core.array.var.rst", "python/_autosummary/mlx.core.array_equal.rst", "python/_autosummary/mlx.core.broadcast_to.rst", "python/_autosummary/mlx.core.concatenate.rst", "python/_autosummary/mlx.core.conv1d.rst", "python/_autosummary/mlx.core.conv2d.rst", "python/_autosummary/mlx.core.convolve.rst", "python/_autosummary/mlx.core.cos.rst", "python/_autosummary/mlx.core.cosh.rst", "python/_autosummary/mlx.core.default_device.rst", "python/_autosummary/mlx.core.default_stream.rst", "python/_autosummary/mlx.core.divide.rst", "python/_autosummary/mlx.core.equal.rst", "python/_autosummary/mlx.core.erf.rst", "python/_autosummary/mlx.core.erfinv.rst", "python/_autosummary/mlx.core.eval.rst", "python/_autosummary/mlx.core.exp.rst", "python/_autosummary/mlx.core.expand_dims.rst", "python/_autosummary/mlx.core.fft.fft.rst", "python/_autosummary/mlx.core.fft.fft2.rst", "python/_autosummary/mlx.core.fft.fftn.rst", "python/_autosummary/mlx.core.fft.ifft.rst", "python/_autosummary/mlx.core.fft.ifft2.rst", "python/_autosummary/mlx.core.fft.ifftn.rst", "python/_autosummary/mlx.core.fft.irfft.rst", "python/_autosummary/mlx.core.fft.irfft2.rst", "python/_autosummary/mlx.core.fft.irfftn.rst", "python/_autosummary/mlx.core.fft.rfft.rst", "python/_autosummary/mlx.core.fft.rfft2.rst", "python/_autosummary/mlx.core.fft.rfftn.rst", "python/_autosummary/mlx.core.full.rst", "python/_autosummary/mlx.core.grad.rst", "python/_autosummary/mlx.core.greater.rst", "python/_autosummary/mlx.core.greater_equal.rst", "python/_autosummary/mlx.core.jvp.rst", "python/_autosummary/mlx.core.less.rst", "python/_autosummary/mlx.core.less_equal.rst", "python/_autosummary/mlx.core.load.rst", "python/_autosummary/mlx.core.log.rst", "python/_autosummary/mlx.core.log10.rst", "python/_autosummary/mlx.core.log1p.rst", "python/_autosummary/mlx.core.log2.rst", "python/_autosummary/mlx.core.logaddexp.rst", "python/_autosummary/mlx.core.logical_not.rst", "python/_autosummary/mlx.core.logsumexp.rst", "python/_autosummary/mlx.core.matmul.rst", "python/_autosummary/mlx.core.max.rst", "python/_autosummary/mlx.core.maximum.rst", "python/_autosummary/mlx.core.mean.rst", "python/_autosummary/mlx.core.min.rst", "python/_autosummary/mlx.core.minimum.rst", "python/_autosummary/mlx.core.multiply.rst", "python/_autosummary/mlx.core.negative.rst", "python/_autosummary/mlx.core.new_stream.rst", "python/_autosummary/mlx.core.ones.rst", "python/_autosummary/mlx.core.ones_like.rst", "python/_autosummary/mlx.core.pad.rst", "python/_autosummary/mlx.core.partition.rst", "python/_autosummary/mlx.core.prod.rst", "python/_autosummary/mlx.core.random.bernoulli.rst", "python/_autosummary/mlx.core.random.categorical.rst", "python/_autosummary/mlx.core.random.gumbel.rst", "python/_autosummary/mlx.core.random.key.rst", "python/_autosummary/mlx.core.random.normal.rst", "python/_autosummary/mlx.core.random.randint.rst", "python/_autosummary/mlx.core.random.seed.rst", "python/_autosummary/mlx.core.random.split.rst", "python/_autosummary/mlx.core.random.truncated_normal.rst", "python/_autosummary/mlx.core.random.uniform.rst", "python/_autosummary/mlx.core.reciprocal.rst", "python/_autosummary/mlx.core.reshape.rst", "python/_autosummary/mlx.core.rsqrt.rst", "python/_autosummary/mlx.core.save.rst", "python/_autosummary/mlx.core.savez.rst", "python/_autosummary/mlx.core.savez_compressed.rst", "python/_autosummary/mlx.core.set_default_device.rst", "python/_autosummary/mlx.core.set_default_stream.rst", "python/_autosummary/mlx.core.sigmoid.rst", "python/_autosummary/mlx.core.sign.rst", "python/_autosummary/mlx.core.sin.rst", "python/_autosummary/mlx.core.sinh.rst", "python/_autosummary/mlx.core.softmax.rst", "python/_autosummary/mlx.core.sort.rst", "python/_autosummary/mlx.core.split.rst", "python/_autosummary/mlx.core.sqrt.rst", "python/_autosummary/mlx.core.square.rst", "python/_autosummary/mlx.core.squeeze.rst", "python/_autosummary/mlx.core.stop_gradient.rst", "python/_autosummary/mlx.core.subtract.rst", "python/_autosummary/mlx.core.sum.rst", "python/_autosummary/mlx.core.take.rst", "python/_autosummary/mlx.core.take_along_axis.rst", "python/_autosummary/mlx.core.tan.rst", "python/_autosummary/mlx.core.tanh.rst", "python/_autosummary/mlx.core.transpose.rst", "python/_autosummary/mlx.core.value_and_grad.rst", "python/_autosummary/mlx.core.var.rst", "python/_autosummary/mlx.core.vjp.rst", "python/_autosummary/mlx.core.vmap.rst", "python/_autosummary/mlx.core.where.rst", "python/_autosummary/mlx.core.zeros.rst", "python/_autosummary/mlx.core.zeros_like.rst", "python/_autosummary/mlx.nn.Conv1d.rst", "python/_autosummary/mlx.nn.Conv2d.rst", "python/_autosummary/mlx.nn.Embedding.rst", "python/_autosummary/mlx.nn.GELU.rst", "python/_autosummary/mlx.nn.GroupNorm.rst", "python/_autosummary/mlx.nn.LayerNorm.rst", "python/_autosummary/mlx.nn.Linear.rst", "python/_autosummary/mlx.nn.MultiHeadAttention.rst", "python/_autosummary/mlx.nn.RMSNorm.rst", "python/_autosummary/mlx.nn.ReLU.rst", "python/_autosummary/mlx.nn.RoPE.rst", "python/_autosummary/mlx.nn.Sequential.rst", "python/_autosummary/mlx.nn.SiLU.rst", "python/_autosummary/mlx.nn.value_and_grad.rst", "python/_autosummary/mlx.optimizers.Adam.rst", "python/_autosummary/mlx.optimizers.Optimizer.rst", "python/_autosummary/mlx.optimizers.OptimizerState.rst", "python/_autosummary/mlx.optimizers.SGD.rst", "python/_autosummary/mlx.utils.tree_flatten.rst", "python/_autosummary/mlx.utils.tree_map.rst", "python/_autosummary/mlx.utils.tree_unflatten.rst", "python/_autosummary_functions/mlx.nn.gelu.rst", "python/_autosummary_functions/mlx.nn.gelu_approx.rst", "python/_autosummary_functions/mlx.nn.gelu_fast_approx.rst", "python/_autosummary_functions/mlx.nn.relu.rst", "python/_autosummary_functions/mlx.nn.silu.rst", "python/array.rst", "python/data_types.rst", "python/devices_and_streams.rst", "python/fft.rst", "python/nn.rst", "python/nn/module.rst", "python/ops.rst", "python/optimizers.rst", "python/random.rst", "python/transforms.rst", "python/tree_utils.rst", "quick_start.rst", "using_streams.rst"], "titles": ["Operations", "Developer Documentation", "Linear Regression", "LLM inference", "Multi-Layer Perceptron", "MLX", "Build and Install", "mlx.core.Device", "mlx.core.Dtype", "mlx.core.Stream", "mlx.core.abs", "mlx.core.add", "mlx.core.all", "mlx.core.allclose", "mlx.core.any", "mlx.core.arange", "mlx.core.arccos", "mlx.core.arccosh", "mlx.core.arcsin", "mlx.core.arcsinh", "mlx.core.arctan", "mlx.core.arctanh", "mlx.core.argmax", "mlx.core.argmin", "mlx.core.argpartition", "mlx.core.argsort", "mlx.core.array", "mlx.core.array.T", "mlx.core.array.abs", "mlx.core.array.all", "mlx.core.array.any", "mlx.core.array.argmax", "mlx.core.array.argmin", "mlx.core.array.astype", "mlx.core.array.cos", "mlx.core.array.dtype", "mlx.core.array.exp", "mlx.core.array.item", "mlx.core.array.log", "mlx.core.array.log1p", "mlx.core.array.logsumexp", "mlx.core.array.max", "mlx.core.array.mean", "mlx.core.array.min", "mlx.core.array.ndim", "mlx.core.array.prod", "mlx.core.array.reciprocal", "mlx.core.array.reshape", "mlx.core.array.rsqrt", "mlx.core.array.shape", "mlx.core.array.sin", "mlx.core.array.size", "mlx.core.array.split", "mlx.core.array.sqrt", "mlx.core.array.square", "mlx.core.array.sum", "mlx.core.array.tolist", "mlx.core.array.transpose", "mlx.core.array.var", "mlx.core.array_equal", "mlx.core.broadcast_to", "mlx.core.concatenate", "mlx.core.conv1d", "mlx.core.conv2d", "mlx.core.convolve", "mlx.core.cos", "mlx.core.cosh", "mlx.core.default_device", "mlx.core.default_stream", "mlx.core.divide", "mlx.core.equal", "mlx.core.erf", "mlx.core.erfinv", "mlx.core.eval", "mlx.core.exp", "mlx.core.expand_dims", "mlx.core.fft.fft", "mlx.core.fft.fft2", "mlx.core.fft.fftn", "mlx.core.fft.ifft", "mlx.core.fft.ifft2", "mlx.core.fft.ifftn", "mlx.core.fft.irfft", "mlx.core.fft.irfft2", "mlx.core.fft.irfftn", "mlx.core.fft.rfft", "mlx.core.fft.rfft2", "mlx.core.fft.rfftn", "mlx.core.full", "mlx.core.grad", "mlx.core.greater", "mlx.core.greater_equal", "mlx.core.jvp", "mlx.core.less", "mlx.core.less_equal", "mlx.core.load", "mlx.core.log", "mlx.core.log10", "mlx.core.log1p", "mlx.core.log2", "mlx.core.logaddexp", "mlx.core.logical_not", "mlx.core.logsumexp", "mlx.core.matmul", "mlx.core.max", "mlx.core.maximum", "mlx.core.mean", "mlx.core.min", "mlx.core.minimum", "mlx.core.multiply", "mlx.core.negative", "mlx.core.new_stream", "mlx.core.ones", "mlx.core.ones_like", "mlx.core.pad", "mlx.core.partition", "mlx.core.prod", "mlx.core.random.bernoulli", "mlx.core.random.categorical", "mlx.core.random.gumbel", "mlx.core.random.key", "mlx.core.random.normal", "mlx.core.random.randint", "mlx.core.random.seed", "mlx.core.random.split", "mlx.core.random.truncated_normal", "mlx.core.random.uniform", "mlx.core.reciprocal", "mlx.core.reshape", "mlx.core.rsqrt", "mlx.core.save", "mlx.core.savez", "mlx.core.savez_compressed", "mlx.core.set_default_device", "mlx.core.set_default_stream", "mlx.core.sigmoid", "mlx.core.sign", "mlx.core.sin", "mlx.core.sinh", "mlx.core.softmax", "mlx.core.sort", "mlx.core.split", "mlx.core.sqrt", "mlx.core.square", "mlx.core.squeeze", "mlx.core.stop_gradient", "mlx.core.subtract", "mlx.core.sum", "mlx.core.take", "mlx.core.take_along_axis", "mlx.core.tan", "mlx.core.tanh", "mlx.core.transpose", "mlx.core.value_and_grad", "mlx.core.var", "mlx.core.vjp", "mlx.core.vmap", "mlx.core.where", "mlx.core.zeros", "mlx.core.zeros_like", "mlx.nn.Conv1d", "mlx.nn.Conv2d", "mlx.nn.Embedding", "mlx.nn.GELU", "mlx.nn.GroupNorm", "mlx.nn.LayerNorm", "mlx.nn.Linear", "mlx.nn.MultiHeadAttention", "mlx.nn.RMSNorm", "mlx.nn.ReLU", "mlx.nn.RoPE", "mlx.nn.Sequential", "mlx.nn.SiLU", "mlx.nn.value_and_grad", "mlx.optimizers.Adam", "mlx.optimizers.Optimizer", "mlx.optimizers.OptimizerState", "mlx.optimizers.SGD", "mlx.utils.tree_flatten", "mlx.utils.tree_map", "mlx.utils.tree_unflatten", "mlx.nn.gelu", "mlx.nn.gelu_approx", "mlx.nn.gelu_fast_approx", "mlx.nn.relu", "mlx.nn.silu", "Array", "Data Types", "Devices and Streams", "FFT", "Neural Networks", "mlx.nn.Module", "Operations", "Optimizers", "Random", "Transforms", "Tree Utils", "Quick Start Guide", "Using Streams"], "terms": {"mlx": [1, 2, 3, 4, 6, 190, 193, 194, 196, 197], "provid": [1, 3, 167, 179, 190, 191, 198], "open": [1, 15, 122, 126], "flexibl": [1, 5], "which": [1, 3, 4, 5, 6, 15, 33, 73, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 89, 92, 95, 118, 119, 128, 130, 131, 132, 144, 148, 153, 155, 156, 163, 170, 191, 194, 198], "user": [1, 3, 190], "mai": 1, "add": [1, 3, 75, 100, 114, 160, 161], "special": 1, "without": [1, 3, 5, 145, 167, 190, 196], "much": [1, 3], "hassl": 1, "while": [1, 3, 6, 128, 170], "librari": [1, 6, 190], "suppli": 1, "effici": [1, 3, 5, 170, 197], "can": [1, 3, 5, 6, 11, 15, 47, 57, 69, 70, 73, 90, 91, 93, 94, 100, 105, 108, 109, 117, 118, 122, 125, 126, 146, 153, 162, 171, 190, 191, 193, 194, 196, 197, 198], "compos": [1, 5, 190, 197], "ani": [1, 3, 5, 15, 163, 178, 179, 180, 190, 191, 196, 197], "number": [1, 15, 51, 63, 89, 92, 114, 118, 121, 124, 126, 153, 155, 156, 160, 161, 164, 167, 194, 198], "applic": 1, "aris": 1, "case": [1, 3, 78, 81, 82, 84, 85, 86, 87, 103, 128, 144, 197, 198], "where": [1, 4, 153, 156, 160, 161, 163, 164, 165, 168, 172, 181, 182, 183, 185, 191], "new": [1, 4, 60, 128, 152, 167, 179, 191, 193], "function": [1, 2, 3, 4, 5, 13, 26, 71, 72, 73, 89, 92, 103, 135, 153, 155, 156, 163, 171, 173, 179, 181, 182, 183, 190, 191, 193, 194, 196], "highli": [1, 6], "optim": [1, 2, 4, 5, 191], "ar": [1, 2, 3, 4, 5, 13, 15, 59, 60, 64, 77, 78, 80, 81, 83, 84, 86, 87, 92, 103, 114, 115, 117, 118, 119, 122, 125, 126, 131, 132, 144, 148, 155, 156, 160, 161, 164, 165, 167, 178, 179, 190, 191, 196, 197], "need": [1, 3, 4, 5, 59, 190, 191, 194, 197], "For": [1, 3, 6, 180, 191, 194, 197], "you": [1, 3, 5, 6, 194], "design": [1, 2, 5, 194], "your": [1, 3, 6, 191], "own": [1, 6], "link": [1, 6], "top": 1, "core": [1, 2, 3, 4, 190, 191, 193, 197], "we": [1, 2, 3, 4, 162, 171, 190, 194, 196], "inner": 1, "work": [1, 3, 6], "go": [1, 3], "over": [1, 3, 4, 12, 14, 22, 23, 24, 25, 62, 63, 78, 81, 84, 87, 102, 104, 106, 107, 115, 116, 129, 139, 140, 147, 154, 160, 161, 164, 165, 168], "simpl": [1, 3, 4, 162, 190], "learn": [1, 2, 4, 5, 164, 165, 168, 177], "step": [1, 3, 4, 15], "involv": [1, 193], "ad": [1, 2, 6, 191], "let": [1, 2, 3], "s": [1, 2, 3, 4, 35, 44, 77, 78, 80, 81, 83, 84, 86, 87, 89, 95, 106, 118, 153, 154, 156, 173, 175, 190, 191, 193, 194, 197], "sai": [1, 3], "would": [1, 3], "like": [1, 3, 5, 113, 159, 197], "an": [1, 3, 4, 6, 8, 12, 14, 26, 60, 62, 63, 73, 88, 104, 107, 112, 113, 114, 116, 128, 141, 144, 148, 149, 156, 158, 159, 164, 165, 166, 167, 175, 176, 178, 182, 190, 191, 194, 196, 197, 198], "take": [1, 3, 4, 89, 92, 105, 108, 113, 149, 153, 155, 156, 159, 194, 198], "two": [1, 11, 13, 59, 69, 70, 77, 80, 86, 90, 91, 93, 94, 100, 103, 105, 108, 109], "arrai": [1, 3, 4, 5, 8, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 59, 60, 61, 62, 63, 64, 65, 66, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 124, 125, 126, 127, 128, 129, 130, 131, 132, 135, 136, 137, 138, 139, 140, 141, 142, 143, 144, 145, 146, 147, 148, 149, 150, 151, 152, 153, 154, 155, 156, 157, 158, 159, 190, 191, 197], "x": [1, 2, 3, 4, 71, 119, 131, 135, 157, 163, 164, 165, 168, 169, 172, 179, 181, 182, 183, 184, 185, 190, 191, 193, 197], "y": [1, 2, 3, 4, 157, 164, 165, 168, 190, 193], "scale": [1, 3, 167], "them": [1, 3, 190, 191], "both": [1, 11, 69, 70, 90, 91, 93, 94, 100, 105, 108, 109, 118, 146, 193, 197], "some": [1, 2, 3, 4, 191], "coeffic": 1, "alpha": 1, "beta": [1, 164, 165, 174], "respect": [1, 2, 4, 89, 153, 163, 164, 165, 179, 190, 191, 197], "togeth": [1, 4, 179], "get": [1, 2, 4, 63, 120, 176], "z": 1, "well": [1, 3, 167, 190, 191], "veri": [1, 3, 167], "easili": 1, "do": [1, 3, 6, 191], "just": [1, 4], "write": [1, 3, 190], "out": [1, 6], "follow": [1, 3, 4, 5, 6, 15, 64, 174, 177, 182, 183, 194], "import": [1, 2, 3, 4, 131, 153, 178, 179, 180, 190, 191, 197], "mx": [1, 2, 3, 4, 131, 153, 169, 184, 190, 191, 193, 194, 197, 198], "def": [1, 2, 3, 4, 153, 190, 191], "simple_axpbi": 1, "float": [1, 11, 13, 15, 26, 56, 59, 60, 69, 70, 88, 90, 91, 93, 94, 100, 101, 105, 108, 109, 114, 117, 122, 125, 126, 146, 157, 164, 165, 168, 174, 177, 187, 191], "return": [1, 2, 3, 4, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 33, 37, 49, 56, 59, 60, 61, 62, 63, 64, 65, 66, 69, 70, 71, 72, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 124, 125, 126, 127, 128, 129, 135, 136, 137, 138, 139, 140, 141, 142, 143, 144, 145, 146, 147, 148, 149, 150, 151, 152, 153, 154, 155, 156, 157, 158, 159, 173, 178, 179, 180, 190, 191, 196], "thi": [1, 3, 4, 12, 13, 14, 15, 22, 23, 24, 25, 73, 89, 92, 100, 102, 103, 104, 106, 107, 115, 116, 118, 139, 140, 141, 147, 148, 153, 154, 182, 183, 190, 191, 196], "perform": [1, 3, 5, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 103, 139, 148, 164, 190], "leav": [1, 179], "differenti": [1, 5], "howev": [1, 163, 164, 190, 194], "vector": [1, 2, 5, 92, 148, 155, 156, 162, 197], "math": [1, 3], "often": 1, "realiz": 1, "axpbi": 1, "routin": 1, "defin": [1, 2, 3, 4, 6, 176, 178], "same": [1, 3, 6, 59, 60, 63, 64, 82, 85, 86, 87, 89, 92, 114, 118, 155, 157, 164, 190, 191, 194], "realli": 1, "part": 1, "doe": [1, 3, 190], "fast": [1, 163, 183], "so": [1, 3, 6, 193], "decid": [1, 191], "want": [1, 3], "reli": 1, "acceler": 1, "framework": [1, 5], "continu": 1, "impos": 1, "our": [1, 3, 171, 174], "assumpt": 1, "also": [1, 3, 4, 5, 11, 69, 70, 78, 81, 84, 87, 90, 91, 93, 94, 100, 105, 108, 109, 146, 167, 173, 176, 181, 190, 191, 193, 197, 198], "assum": [1, 3, 164, 179, 190], "how": [1, 3, 4, 160, 161, 162, 167, 190], "gradient": [1, 2, 4, 89, 145, 153, 173, 177, 190, 191, 193, 197], "ins": 1, "what": [1, 3], "coincid": 1, "right": [1, 163, 182, 183], "place": [1, 3], "cours": 1, "The": [1, 3, 4, 5, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 33, 35, 44, 49, 56, 59, 60, 61, 62, 63, 64, 65, 66, 69, 70, 71, 72, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 124, 125, 126, 127, 128, 135, 136, 137, 138, 139, 140, 142, 143, 144, 145, 146, 147, 148, 149, 150, 151, 152, 153, 154, 155, 156, 157, 158, 159, 160, 161, 162, 164, 165, 166, 167, 168, 170, 171, 173, 175, 176, 177, 178, 179, 180, 187, 191, 193, 197, 198], "structur": [1, 73], "from": [1, 3, 4, 5, 83, 84, 86, 87, 88, 95, 103, 113, 117, 118, 119, 120, 122, 125, 131, 144, 145, 146, 148, 149, 157, 159, 167, 178, 179, 180, 190, 191, 196, 197], "frontend": 1, "api": 1, "redirect": 1, "when": [1, 3, 5, 160, 161, 190, 191, 194], "appropri": 1, "fallback": 1, "metal": [1, 6], "vjp": [1, 197], "jvp": [1, 197], "In": [1, 3, 4, 103, 164, 174, 179, 190, 191, 196], "one": [1, 3, 6, 56, 63, 75, 98, 103, 118, 144, 146, 191], "sentenc": 1, "comput": [1, 2, 3, 4, 5, 6, 89, 92, 100, 106, 139, 145, 153, 154, 155, 164, 165, 168, 173, 182, 183, 190, 191, 193, 197], "graph": [1, 3, 4, 5, 73, 130], "rule": 1, "evalu": [1, 3, 4, 73, 92, 130, 155, 190, 191, 193, 197], "said": [1, 3], "start": [1, 2, 3, 5, 6, 15, 141], "discuss": 1, "more": [1, 4, 8, 56, 103, 190, 194], "detail": [1, 8, 174, 190], "thei": [1, 2, 3, 64, 171, 191, 196, 197], "c": [1, 3, 160, 161, 187, 197], "scalar": [1, 11, 13, 37, 56, 59, 69, 70, 89, 90, 91, 93, 94, 100, 101, 103, 105, 108, 109, 114, 122, 125, 126, 146, 153, 173, 197], "sum": [1, 2, 11, 102, 139], "elementwis": 1, "numpi": [1, 3, 4, 5, 11, 13, 15, 26, 60, 69, 70, 90, 91, 93, 94, 100, 103, 105, 108, 109, 146, 197], "style": [1, 11, 13, 69, 70, 90, 91, 93, 94, 100, 103, 105, 108, 109, 146], "broadcast": [1, 11, 13, 60, 69, 70, 88, 90, 91, 93, 94, 100, 103, 105, 108, 109, 117, 118, 125, 126, 146, 149, 157, 167], "between": [1, 5], "input": [1, 2, 3, 10, 11, 12, 13, 14, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 59, 60, 61, 62, 63, 64, 65, 66, 69, 70, 71, 72, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 89, 90, 91, 92, 93, 94, 96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 113, 114, 115, 116, 124, 127, 128, 129, 135, 136, 137, 138, 139, 140, 141, 142, 143, 144, 145, 146, 147, 148, 149, 150, 151, 152, 153, 154, 156, 157, 159, 160, 161, 162, 164, 165, 166, 167, 168, 170, 197], "upcast": 1, "const": 1, "factor": 1, "streamordevic": 1, "stream": [1, 5, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 28, 29, 30, 31, 32, 33, 34, 36, 38, 39, 40, 41, 42, 43, 45, 46, 47, 48, 50, 52, 53, 54, 55, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 68, 69, 70, 71, 72, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 90, 91, 93, 94, 95, 96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115, 116, 117, 118, 119, 121, 122, 124, 125, 126, 127, 128, 129, 134, 135, 136, 137, 138, 139, 140, 141, 142, 143, 144, 145, 146, 147, 148, 149, 150, 151, 152, 154, 157, 158, 159], "schedul": 1, "itself": 1, "call": [1, 3, 4, 27, 162, 171, 190, 191, 193], "other": [1, 3, 167, 190, 191, 197], "within": [1, 24], "simplest": 1, "wai": [1, 3, 190], "about": [1, 3, 4], "term": 1, "exist": [1, 3, 191], "auto": 1, "ax": [1, 12, 14, 22, 23, 57, 75, 77, 78, 80, 81, 83, 84, 86, 87, 102, 104, 106, 107, 114, 116, 139, 144, 147, 152, 154], "multipli": 1, "earlier": 1, "goal": 1, "themselv": 1, "contain": [1, 3, 49, 73, 82, 83, 84, 101, 141, 157, 190, 191], "act": 1, "data": [1, 4, 5, 8, 15, 85, 86, 88, 112, 125, 158], "nor": 1, "rather": 1, "easi": [1, 190], "interfac": 1, "block": [1, 3], "A": [1, 3, 5, 6, 49, 59, 89, 92, 102, 103, 117, 118, 119, 121, 122, 125, 126, 141, 153, 155, 156, 164, 165, 168, 171, 173, 174, 178, 179, 180, 183, 190, 191, 193], "It": [1, 3, 6, 175, 190], "creat": [1, 3, 6, 190, 191, 193], "output": [1, 3, 12, 13, 14, 15, 22, 23, 24, 60, 82, 85, 86, 87, 88, 89, 102, 104, 106, 107, 112, 113, 115, 116, 117, 118, 119, 121, 122, 125, 126, 131, 132, 139, 144, 147, 149, 153, 154, 155, 156, 157, 158, 159, 160, 161, 166, 167, 197], "given": [1, 12, 14, 24, 60, 61, 73, 75, 76, 77, 78, 79, 80, 81, 85, 86, 87, 88, 102, 104, 106, 107, 116, 122, 139, 141, 147, 154, 167, 191], "set": [1, 3, 4, 163, 166, 170, 176, 191, 194], "further": [1, 6], "class": [1, 3, 4, 7, 8, 9, 26, 160, 161, 162, 163, 164, 165, 166, 167, 168, 169, 170, 171, 172, 174, 175, 176, 177, 181, 182, 183, 184, 185, 191], "under": 1, "These": [1, 149], "word": 1, "bit": [1, 187, 191], "abstract": 1, "back": [1, 3], "give": [1, 3, 4, 24], "ourselv": 1, "concret": 1, "imag": [1, 161], "public": [1, 190], "explicit": [1, 194], "alpha_": 1, "beta_": 1, "must": [1, 6, 73, 88, 117, 118, 122, 125, 126, 157], "know": [1, 3], "popul": 1, "To": [1, 2, 3, 6, 197], "avoid": 1, "unecessari": [], "alloc": [1, 191], "respons": 1, "space": 1, "void": 1, "eval_cpu": 1, "std": 1, "overrid": 1, "eval_gpu": 1, "jacobian": [1, 92, 155, 197], "product": [1, 92, 103, 116, 155, 167, 197], "primal": [1, 92, 155], "tangent": [1, 20, 21, 92, 150, 151], "int": [1, 3, 4, 7, 9, 11, 12, 14, 15, 22, 23, 24, 25, 26, 29, 30, 31, 32, 40, 41, 42, 43, 45, 49, 52, 55, 56, 58, 59, 60, 61, 62, 63, 69, 70, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 93, 94, 100, 101, 102, 104, 105, 106, 107, 108, 109, 112, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 128, 139, 140, 141, 144, 146, 147, 148, 149, 152, 153, 154, 156, 157, 158, 160, 161, 162, 164, 165, 166, 167, 168, 170, 190, 191], "argnum": [1, 89, 153], "cotan": 1, "accross": [1, 164], "pair": [1, 114, 170], "repres": [1, 3], "axi": [1, 3, 4, 12, 14, 22, 23, 24, 25, 29, 30, 31, 32, 40, 41, 42, 43, 45, 52, 55, 58, 61, 75, 76, 79, 82, 83, 84, 85, 86, 87, 102, 104, 106, 107, 114, 115, 116, 118, 139, 140, 141, 144, 147, 148, 149, 152, 154, 156], "correspond": [1, 12, 14, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 102, 104, 107, 116, 147, 156, 179], "dimens": [1, 3, 12, 14, 22, 23, 44, 49, 56, 63, 75, 83, 84, 86, 87, 102, 103, 104, 106, 107, 116, 118, 124, 147, 149, 152, 154, 160, 161, 164, 165, 167, 168, 170], "vmap": [1, 197], "print": [1, 2, 3, 4, 178, 179, 180, 190, 194, 197], "ostream": 1, "os": 1, "equival": [1, 27, 47, 57, 163], "check": [1, 59], "bool": [1, 11, 12, 14, 22, 23, 26, 29, 30, 31, 32, 40, 41, 42, 43, 45, 55, 56, 58, 59, 60, 69, 70, 73, 88, 90, 91, 93, 94, 100, 101, 102, 104, 105, 106, 107, 108, 109, 114, 116, 117, 122, 125, 126, 130, 146, 147, 154, 157, 160, 161, 164, 165, 166, 170, 191], "is_equival": 1, "privat": 1, "fall": 1, "eval": [1, 2, 3, 4, 190, 191, 193, 197], "deriv": 1, "base": [1, 73, 97, 99, 175, 191, 193, 194], "abov": [1, 3, 6], "demonstr": 1, "treat": [1, 59, 83, 84, 86, 87, 148], "paramet": [1, 2, 3, 4, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 33, 59, 60, 61, 62, 63, 64, 65, 66, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 127, 128, 129, 130, 131, 132, 135, 136, 137, 138, 139, 140, 141, 142, 143, 144, 145, 146, 147, 148, 149, 150, 151, 152, 153, 154, 155, 156, 157, 158, 159, 160, 161, 162, 163, 164, 165, 166, 167, 168, 170, 171, 173, 175, 177, 178, 179, 180, 191, 193], "produc": [1, 167], "through": [1, 145], "construct": [1, 4, 88, 112, 158], "its": [1, 6, 103, 115, 124, 173, 180, 190], "type": [1, 5, 7, 8, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 33, 49, 56, 59, 60, 61, 62, 63, 64, 65, 66, 69, 70, 71, 72, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 124, 125, 126, 127, 128, 129, 135, 136, 137, 138, 139, 140, 141, 142, 143, 144, 145, 146, 147, 148, 149, 150, 151, 152, 153, 154, 155, 156, 157, 158, 159, 175, 178, 190], "shape": [1, 3, 4, 47, 59, 60, 62, 63, 76, 79, 82, 85, 86, 87, 88, 92, 103, 112, 113, 117, 118, 119, 121, 122, 125, 126, 128, 149, 155, 157, 158, 159, 160, 161, 190, 193, 197], "pass": [1, 3, 4, 47, 57, 114, 153, 171, 173, 178, 190, 191], "re": [1, 4], "now": [1, 3], "promot": 1, "dtype": [1, 3, 15, 26, 33, 56, 88, 112, 119, 121, 122, 125, 126, 158, 187, 197], "promoted_dtyp": 1, "promote_typ": 1, "float32": [1, 15, 112, 119, 121, 125, 126, 158, 187, 197], "non": [1, 191], "point": [1, 2, 3, 6, 187], "out_dtyp": 1, "is_floating_point": 1, "cast": [1, 33, 85, 86, 87, 191], "up": [1, 3], "determin": 1, "x_cast": 1, "astyp": [1, 3, 191], "y_cast": 1, "broadcasted_input": 1, "broadcast_arrai": 1, "out_shap": 1, "0": [1, 2, 3, 4, 6, 7, 15, 52, 58, 61, 62, 63, 89, 114, 117, 126, 141, 153, 154, 156, 160, 161, 163, 164, 165, 169, 174, 177, 178, 182, 183, 184, 190, 191, 194, 197], "unique_ptr": 1, "make_uniqu": 1, "to_stream": 1, "handl": [1, 190], "resolv": 1, "No": [1, 3], "happen": [1, 3, 193], "alon": 1, "effect": 1, "onli": [1, 3, 5, 6, 59, 62, 63, 187, 190, 191], "execut": [1, 6], "depend": [1, 2, 56], "devic": [1, 5, 6, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 28, 29, 30, 31, 32, 33, 34, 36, 38, 39, 40, 41, 42, 43, 45, 46, 47, 48, 50, 52, 53, 54, 55, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 90, 91, 93, 94, 95, 96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115, 116, 117, 118, 119, 121, 122, 124, 125, 126, 127, 128, 129, 133, 135, 136, 137, 138, 139, 140, 141, 142, 143, 144, 145, 146, 147, 148, 149, 150, 151, 152, 154, 157, 158, 159, 198], "specifi": [1, 15, 33, 63, 83, 84, 88, 89, 112, 118, 148, 149, 152, 153, 156, 158], "memori": [1, 5, 191], "ha": [1, 3, 4, 5, 56, 82, 83, 85, 86, 87, 89, 118, 191, 193, 197], "been": [1, 3], "try": 1, "naiv": 1, "gener": [1, 2, 15, 83, 84, 117, 121, 122, 125, 126, 194, 198], "version": [1, 6, 100, 102, 139, 156, 194], "declar": 1, "member": [1, 190, 191], "method": [1, 3, 7, 8, 9, 26, 174, 175, 176, 177], "each": [1, 49, 73, 103, 114, 118, 131, 132, 141, 152, 156, 157, 162, 164, 194], "element": [1, 10, 11, 16, 17, 18, 19, 20, 21, 24, 65, 66, 69, 70, 71, 72, 74, 90, 91, 93, 94, 96, 97, 98, 99, 100, 101, 105, 108, 109, 110, 115, 127, 129, 135, 136, 137, 138, 142, 143, 146, 148, 150, 151, 153, 157, 170, 172, 185], "find": [1, 2, 6], "pointwis": 1, "captur": [1, 190], "templat": 1, "axpby_impl": 1, "typenam": 1, "t": [1, 3, 71, 153, 174, 177, 190], "readi": 1, "fill": [1, 88, 113, 159], "malloc_or_wait": 1, "synchron": 1, "avail": [1, 2, 3, 4, 6, 8, 187], "There": [1, 190], "wait": [1, 3], "here": [1, 3], "request": 1, "pressur": 1, "condit": [1, 157], "set_data": 1, "nbyte": 1, "collect": [1, 176, 179, 196], "pointer": 1, "x_ptr": 1, "y_ptr": 1, "out_ptr": 1, "relev": 1, "static_cast": 1, "size_t": 1, "out_idx": 1, "size": [1, 3, 4, 49, 63, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 118, 128, 141, 144, 160, 161, 162], "map": [1, 4, 95, 162, 179, 191], "linear": [1, 3, 4, 5, 163, 169, 172, 179, 181, 182, 183, 184, 185, 190, 191], "indic": [1, 13, 22, 23, 24, 25, 73, 89, 141, 148, 149, 153], "offset": [1, 3], "x_offset": 1, "elem_to_loc": 1, "stride": [1, 62, 63, 160, 161, 170], "y_offset": 1, "contigu": 1, "regularli": 1, "default": [1, 6, 12, 14, 15, 22, 23, 24, 25, 59, 61, 62, 63, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 89, 102, 104, 106, 107, 112, 115, 116, 117, 118, 119, 121, 122, 124, 125, 126, 128, 130, 140, 141, 144, 147, 152, 153, 154, 156, 158, 160, 161, 167, 170, 176, 187, 191, 194, 196, 198], "row": 1, "major": 1, "henc": 1, "doesn": [1, 190], "additon": 1, "abl": 1, "all": [1, 4, 6, 13, 24, 63, 73, 78, 81, 84, 87, 103, 114, 115, 144, 167, 175, 190, 191, 194, 197, 198], "incom": 1, "accordingli": 1, "dispatch": 1, "float16": [1, 187, 191], "bfloat16": 1, "complex64": 1, "throw": 1, "error": [1, 71, 72, 141, 163, 181, 182, 183], "encount": 1, "unexpect": [1, 15], "regist": [1, 4], "op": 1, "contruct": 1, "assert": 1, "2": [1, 2, 3, 4, 63, 71, 77, 80, 82, 83, 84, 85, 86, 87, 99, 103, 124, 161, 163, 168, 174, 182, 187, 190, 191, 197], "1": [1, 3, 4, 15, 24, 25, 62, 63, 76, 77, 79, 80, 82, 83, 84, 85, 86, 87, 103, 115, 118, 126, 135, 140, 148, 153, 160, 161, 163, 164, 165, 168, 170, 174, 177, 182, 183, 187, 191, 193, 197], "correct": [1, 174], "els": [1, 3, 190, 191], "float16_t": 1, "bfloat16_t": 1, "complex64_t": 1, "runtime_error": 1, "support": [1, 3, 5, 6, 13, 62, 63, 103], "have": [1, 3, 6, 59, 83, 84, 86, 87, 103, 118, 167, 171, 178, 196], "rememb": 1, "3": [1, 3, 6, 194, 197], "complic": 1, "keep": [1, 12, 14, 22, 23, 102, 104, 106, 107, 116, 147, 154, 190, 191], "mind": [1, 3], "half": [1, 15, 122, 126, 170], "precis": [1, 3, 163, 190], "direct": [1, 3, 191], "fix": [1, 3], "possibl": [1, 3, 103, 141, 162], "due": 1, "transpos": [1, 3, 27], "aren": 1, "guarante": 1, "fit": 1, "requir": [1, 3, 190], "column": 1, "inplac": 1, "expect": [1, 3, 160, 161, 167], "answer": 1, "copi": [1, 3, 5, 115, 140], "simpli": [1, 3, 6, 169, 184, 191], "catlas_saxpbi": 1, "axpby_impl_acceler": 1, "first": [1, 2, 3, 4, 6, 89, 103, 115, 124, 153, 164, 174, 178, 190], "mode": [1, 64], "i": [1, 3, 92, 160, 161, 190], "e": [1, 4, 6, 71, 92, 135, 160, 161, 164, 165, 168, 190, 193, 198], "match": [1, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87], "transposit": 1, "data_s": 1, "items": 1, "flag": 1, "copy_inplac": 1, "copytyp": 1, "n": [1, 3, 26, 62, 63, 76, 78, 79, 81, 82, 85, 87, 154, 160, 161], "incx": 1, "inci": 1, "great": 1, "But": 1, "criteria": 1, "luckili": 1, "alwai": [1, 178], "With": 1, "final": [1, 2, 3, 4], "singl": [1, 4, 73, 92, 114, 155], "row_contigu": 1, "col_contigu": 1, "common": 1, "hit": 1, "mileston": 1, "enough": 1, "run": [1, 3, 4, 5, 6, 191, 198], "If": [1, 3, 6, 12, 14, 15, 22, 23, 24, 25, 56, 59, 61, 64, 73, 85, 86, 87, 88, 89, 102, 103, 104, 106, 107, 112, 114, 115, 116, 118, 139, 140, 141, 147, 148, 149, 153, 154, 156, 158, 160, 161, 164, 165, 166, 167, 170, 171, 179, 191, 198], "plan": 1, "stop": [1, 3, 15, 145], "enjoi": 1, "speed": 1, "appl": [1, 3, 5, 6], "silicon": [1, 3, 5, 6], "address": 1, "shade": 1, "languag": [1, 187], "kernel": [1, 62, 63], "written": 1, "help": [1, 3], "resourc": 1, "walkthrough": 1, "pipelin": 1, "specif": [1, 6], "cpp": 1, "algorithm": 1, "launch": 1, "exactli": [1, 3], "mani": [1, 141, 160, 161, 162, 167], "thread": 1, "pick": 1, "updat": [1, 2, 3, 4, 177, 179, 191, 193], "assign": [1, 191], "axpby_gener": 1, "buffer": [1, 26], "constant": [1, 3, 6, 114, 164, 165, 168], "4": [1, 3, 131, 187, 197], "5": [1, 2, 3, 6, 117], "x_stride": 1, "6": [1, 3, 131, 182, 183, 197], "y_stride": 1, "7": [1, 3], "ndim": 1, "8": [1, 3, 187, 197], "uint": 1, "index": [1, 7, 9, 24, 75, 89, 115, 148, 149, 153], "thread_position_in_grid": 1, "convert": [1, 56, 197], "instanti": [1, 4], "uniqu": [1, 194], "host": 1, "name": [1, 95, 131, 132, 164, 176, 190, 191], "identifi": [1, 178, 196], "instantiate_axpbi": 1, "type_nam": 1, "host_nam": 1, "axpby_general_": 1, "bflot16": 1, "compil": [1, 6], "mlx_ext": 1, "metallib": [1, 6], "see": [1, 3, 4, 8, 28, 29, 30, 31, 32, 34, 36, 38, 39, 40, 41, 42, 43, 45, 46, 47, 48, 50, 52, 53, 54, 55, 57, 58, 163, 181, 182, 183], "later": [1, 6], "co": 1, "locat": [1, 191], "share": [1, 5], "register_librari": 1, "potenti": 1, "path": [1, 6, 131, 132], "tri": 1, "load": [1, 4, 191], "hasn": 1, "alreadi": [1, 3], "static": [1, 6], "object": [1, 8, 26, 37, 56, 95, 130, 131, 132, 156, 178, 196], "why": [1, 3], "packag": [1, 2, 4], "process": [1, 3, 64, 162, 179, 196], "logic": [1, 101], "grid": 1, "shown": 1, "below": [1, 187], "prepar": [1, 3], "carri": 1, "should": [1, 2, 3, 4, 6, 73, 92, 149, 153, 155, 160, 161, 167, 171, 178, 190, 191, 196, 198], "d": [1, 3, 103, 148, 174, 180], "ostringstream": 1, "kname": 1, "axpby_": 1, "general_": 1, "type_to_nam": 1, "make": [1, 3, 4, 6, 103, 190, 197], "sure": [1, 3, 6, 190], "look": [1, 3], "folder": 1, "get_colocated_mtllib_path": 1, "get_kernel": 1, "str": [1, 64, 95, 130, 131, 132, 178, 180, 191], "encod": [1, 170], "compute_encod": 1, "get_command_encod": 1, "setcomputepipelinest": 1, "those": [1, 3, 190], "decelar": 1, "nelem": 1, "set_array_buff": 1, "setbyt": 1, "sizeof": 1, "threadgroup": 1, "higher": 1, "than": [1, 3, 56, 64, 90, 91, 93, 94, 103, 170, 179], "max": [1, 105], "allow": [1, 175, 190, 191, 197], "tgp_size": 1, "min": [1, 108], "maxtotalthreadsperthreadgroup": 1, "3d": 1, "mtl": 1, "group_dim": 1, "grid_dim": 1, "divd": 1, "among": 1, "dispatchthread": 1, "few": [1, 3, 4, 5, 197], "thing": [1, 3], "note": [1, 3, 6, 13, 62, 63, 83, 84, 118, 190], "befor": [1, 3, 6, 24, 115, 130, 191], "move": 1, "track": [1, 190], "activ": [1, 190], "command": [1, 6], "instead": [1, 190], "end_encod": 1, "end": 1, "until": [1, 197], "limit": 1, "flush": 1, "enqueu": 1, "commit": 1, "associ": [1, 131, 132], "suggest": 1, "deeper": 1, "dive": 1, "studi": 1, "come": [1, 3], "far": [1, 193], "built": [1, 6], "includ": [1, 191, 197, 198], "forward": [1, 153], "diff": 1, "push": 1, "along": [1, 22, 23, 61, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 139, 141, 148, 149], "primtiv": 1, "similarli": [1, 6, 103], "scale_arr": 1, "contribut": 1, "tangent_x": 1, "tangent_i": 1, "revers": [1, 152], "arg": [1, 3, 8, 26, 47, 57, 73, 131, 132], "push_back": 1, "fulli": [1, 5], "primitv": 1, "overal": 1, "directori": [1, 3, 6], "extens": [1, 187], "h": [1, 62, 63, 161], "mlx_sample_extens": 1, "__init__": [1, 3, 4, 7, 8, 9, 26, 190, 191], "py": [1, 3], "cmakelist": 1, "txt": 1, "setup": [1, 2, 4], "strucutr": 1, "hold": [1, 3, 8, 175], "instal": 1, "pybind11": [1, 6], "sinc": [1, 3, 4, 191], "compon": [1, 3], "etc": [1, 190], "becom": 1, "pybind11_modul": 1, "m": [1, 6], "doc": [1, 4], "sampl": [1, 2, 3, 117, 118, 119, 122, 125, 126, 194], "_a": 1, "pos_onli": 1, "kw_onli": 1, "none": [1, 3, 7, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 28, 29, 30, 31, 32, 33, 34, 36, 38, 39, 40, 41, 42, 43, 45, 46, 47, 48, 50, 52, 53, 54, 55, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 90, 91, 93, 94, 95, 96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 112, 113, 114, 115, 116, 117, 118, 119, 121, 122, 123, 124, 125, 126, 127, 128, 129, 130, 131, 132, 133, 134, 135, 136, 137, 138, 139, 140, 141, 142, 143, 144, 145, 146, 147, 148, 149, 150, 151, 152, 154, 156, 157, 158, 159, 163, 167, 178, 191], "r": [1, 3, 153], "pbdoc": 1, "most": [1, 118, 190], "complex": [1, 11, 26, 59, 60, 69, 70, 83, 84, 85, 86, 87, 88, 90, 91, 93, 94, 100, 101, 105, 108, 109, 114, 117, 122, 125, 126, 146, 157, 178, 190], "addit": [1, 3, 11, 164, 165, 167, 168, 191], "bell": 1, "whistl": 1, "liter": 1, "string": 1, "modul": [1, 3, 4, 171, 173, 196], "ensur": 1, "caster": 1, "find_packag": 1, "config": 1, "add_librari": 1, "sourc": [1, 152], "target_sourc": 1, "cmake_current_list_dir": 1, "header": 1, "target_include_directori": 1, "target_link_librari": 1, "attach": 1, "conveni": [1, 4], "mlx_build_metallib": 1, "target": [1, 153], "destin": 1, "automat": [1, 5, 197], "practic": 1, "mlx_build_met": [1, 6], "mlx_ext_metallib": 1, "titl": 1, "include_dir": 1, "project_source_dir": 1, "mlx_include_dir": 1, "output_directori": 1, "cmake_library_output_directori": 1, "add_depend": 1, "endif": 1, "pybind11_add_modul": 1, "build_shared_lib": 1, "target_link_opt": 1, "wl": 1, "rpath": 1, "loader_path": 1, "onc": 1, "describ": 1, "util": [1, 3, 5, 131], "__name__": [1, 3], "__main__": [1, 3], "descript": [1, 3, 187], "ext_modul": 1, "cmakeextens": 1, "cmdclass": 1, "build_ext": 1, "cmakebuild": 1, "package_dir": 1, "package_data": 1, "dylib": 1, "zip_saf": 1, "fals": [1, 3, 12, 14, 22, 23, 29, 30, 31, 32, 40, 41, 42, 43, 45, 55, 58, 59, 73, 102, 104, 106, 107, 116, 147, 154, 157, 164, 166, 170, 178, 187, 191], "python_requir": 1, "even": [1, 3], "though": [1, 3], "j8": 1, "libmlx_ext": 1, "cpython": 1, "3x": 1, "darwin": 1, "pip": [1, 6], "after": [1, 3, 4, 24, 115, 164, 165, 167], "plai": [1, 3], "ones": [1, 3, 113, 131, 191], "b": [1, 3, 11, 13, 59, 69, 70, 90, 91, 93, 94, 100, 103, 105, 108, 109, 146, 153, 197], "f": [1, 2, 4, 190], "item": [1, 2, 3, 4, 179, 197], "true": [1, 2, 3, 59, 130, 139, 157, 160, 161, 164, 165, 166, 170, 178, 187, 191], "quick": [1, 5], "benchmark": 1, "compar": [1, 59], "time": [1, 3, 6, 190], "set_default_devic": 1, "256": [1, 4], "512": [1, 3], "random": [1, 2, 3, 4, 5, 198], "normal": [1, 2, 3, 125, 164, 165, 168, 176], "bench": 1, "warm": 1, "rang": [1, 2, 3, 4, 15, 182, 183, 193, 194], "100": [1, 2, 3], "5000": 1, "simple_tim": 1, "custom_tim": 1, "3f": [1, 4], "custom": 1, "114": 1, "109": 1, "modest": 1, "improv": [1, 3], "awai": [1, 3], "good": 1, "nn": [1, 3, 4, 131, 179, 190, 193], "grad": [1, 2, 4, 153, 193, 197], "simplifi": 1, "full": [1, 4, 47, 57, 64, 139, 191], "implement": [2, 4, 162, 167, 170, 171, 174, 175, 176, 191], "basic": 2, "model": [2, 4, 5, 131, 167, 173, 179, 190, 191, 193], "problem": [2, 4, 190], "metadata": 2, "num_featur": 2, "num_exampl": 2, "1_000": 2, "num_it": 2, "10_000": 2, "iter": [2, 4, 179, 194], "sgd": [2, 4, 193], "lr": 2, "01": 2, "rate": 2, "ll": [2, 4], "synthet": 2, "dataset": 2, "matrix": [2, 103], "ground": [2, 3], "truth": 2, "w_star": 2, "valu": [2, 3, 10, 15, 22, 23, 37, 56, 59, 73, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 114, 117, 118, 119, 121, 122, 125, 126, 148, 149, 153, 156, 167, 173, 176, 178, 179, 187, 191], "gaussian": [2, 163, 181, 182, 183], "nois": 2, "exampl": [2, 3, 4, 15, 148, 193, 194, 197], "noisi": 2, "label": 2, "ep": [2, 164, 165, 168, 174], "1e": [2, 4, 13, 164, 165, 168, 174], "us": [2, 3, 4, 5, 6, 15, 103, 128, 162, 163, 166, 167, 175, 178, 182, 183, 190, 191, 193, 194, 196, 197], "weight": [2, 62, 63, 179, 190, 191], "squar": [2, 3, 129, 142, 153, 168, 179, 190], "loss": [2, 4, 153, 190, 193], "loss_fn": [2, 4, 193], "w": [2, 63, 153, 161, 177], "mean": [2, 3, 4, 153, 164, 168, 190, 191], "grad_fn": 2, "initi": [2, 3, 164, 165, 168, 190, 191], "randomli": [2, 3], "Then": [2, 6], "repeatedli": 2, "_": [2, 3, 190, 194], "verifi": 2, "close": [2, 5, 13], "error_norm": 2, "5f": 2, "someth": [2, 3], "00005": 2, "00364": 2, "complet": [2, 3, 191], "logist": [2, 135, 172, 182, 183, 185], "github": [2, 4, 6], "repo": [2, 4, 6], "enabl": [3, 73], "larg": [3, 190], "ish": 3, "transform": [3, 5, 73, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 164, 165, 166, 173, 190, 191], "compromis": 3, "eas": 3, "llama": 3, "famili": 3, "less": [3, 24, 94, 115, 170], "200": 3, "line": 3, "python": [3, 37, 49, 56, 73, 178, 179, 180, 191, 196], "neural": [3, 5, 162, 191], "network": [3, 5, 162, 191], "build": [3, 5, 191], "concis": 3, "architectur": 3, "notabl": 3, "rope": 3, "posit": [3, 24, 115, 160, 161, 167, 170, 179, 190], "option": [3, 12, 14, 15, 22, 23, 24, 25, 26, 31, 32, 61, 62, 63, 64, 73, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 102, 104, 106, 107, 112, 114, 115, 116, 117, 118, 119, 121, 122, 124, 125, 126, 128, 130, 139, 140, 141, 144, 147, 148, 149, 152, 153, 154, 156, 158, 160, 161, 167, 178, 191, 194, 198], "kei": [3, 117, 118, 119, 121, 122, 124, 125, 126, 167, 176, 178, 179, 191, 194, 196], "cach": 3, "concaten": 3, "project": [3, 167], "llamaattent": 3, "self": [3, 4, 7, 9, 26, 27, 28, 29, 30, 31, 32, 33, 34, 36, 37, 38, 39, 40, 41, 42, 43, 45, 46, 47, 48, 50, 52, 53, 54, 55, 56, 57, 58, 190, 191], "dim": [3, 162, 164, 165, 167, 168, 170], "num_head": [3, 167], "super": [3, 4, 190, 191], "tradit": [3, 170], "query_proj": 3, "bia": [3, 160, 161, 166, 174, 179, 191], "key_proj": 3, "value_proj": 3, "out_proj": [3, 191], "__call__": [3, 4, 190, 191], "queri": [3, 167], "mask": [3, 167], "extract": [3, 190, 191], "l": [3, 4, 160, 190], "reshap": 3, "combin": 3, "key_cach": 3, "value_cach": 3, "sqrt": [3, 71, 164, 165, 168, 174], "score": 3, "softmax": 3, "values_hat": 3, "rm": 3, "swiglu": 3, "rmsnorm": 3, "llamaencoderlay": 3, "mlp_dim": 3, "norm1": 3, "norm2": 3, "linear1": 3, "linear2": 3, "linear3": 3, "sigmoid": [3, 172, 182, 183, 185], "instanc": [3, 171, 180, 190, 191], "embed": 3, "emb": [3, 162], "token": [3, 162], "num_lay": [3, 4, 193], "vocab_s": 3, "norm": [3, 164], "multiheadattent": 3, "create_additive_causal_mask": 3, "list": [3, 8, 12, 14, 26, 29, 30, 40, 41, 42, 43, 45, 49, 52, 55, 56, 58, 60, 61, 73, 75, 77, 78, 80, 81, 83, 84, 86, 87, 88, 89, 92, 102, 104, 106, 107, 112, 114, 116, 117, 118, 119, 121, 122, 125, 126, 128, 139, 141, 144, 147, 152, 153, 154, 155, 158, 174, 178, 180, 190, 191, 196], "still": 3, "consid": [3, 13, 59, 164, 178, 196], "train": [3, 4, 191], "ignor": 3, "whatsoev": 3, "rest": [3, 170, 179], "subsect": 3, "prompt": 3, "autoregress": 3, "yield": [3, 4, 194], "temp": 3, "causal": 3, "save": [3, 95, 131, 132, 191], "append": [3, 103], "store": 3, "per": [3, 4, 164, 165, 168, 175], "care": 3, "last": [3, 25, 56, 78, 81, 83, 84, 86, 87, 103, 118, 140, 160, 161, 164], "logit": [3, 118], "next": [3, 4], "categor": 3, "lazili": [3, 190], "noth": [3, 190], "yet": [3, 190, 191, 197], "forc": [3, 4, 190, 197], "choos": [3, 170], "pars": 3, "feed": 3, "loop": [3, 4], "unsqueez": 3, "sequenc": [3, 160, 194], "length": [3, 144, 160], "len": [3, 78, 81, 84, 87], "overwrit": 3, "discard": [3, 178], "old": 3, "moment": [3, 174], "anymor": 3, "everyth": 3, "small": [3, 164, 165, 168], "10": [3, 4, 97, 131, 179, 190], "12": 3, "8192": 3, "1024": 3, "actual": [3, 15, 191], "materi": [3, 5], "could": [3, 190], "20_000": 3, "machin": [3, 5], "8gb": 3, "ram": 3, "32": [3, 4, 187], "44": 3, "doubl": 3, "bracket": 3, "becaus": [3, 190], "batch": [3, 103, 160, 161, 167], "zip": [3, 4], "haven": 3, "anyth": [3, 153], "result": [3, 15, 56, 95, 103, 157, 179], "similar": [3, 167, 191], "runtim": 3, "section": [3, 141], "access": [3, 37, 190, 191], "origin": [3, 174], "sentencepiec": 3, "pytorch": [3, 5, 164], "compat": [3, 118], "npz": [3, 95, 131, 132, 191], "file": [3, 6, 95, 130, 131, 132, 191], "directli": 3, "argpars": 3, "itertool": [3, 179], "starmap": [3, 179], "np": [3, 4, 197], "torch": 3, "map_torch_to_mlx": 3, "tok_embed": 3, "elif": 3, "replac": [3, 191], "attention_norm": 3, "ffn_norm": 3, "wq": 3, "wk": 3, "wv": 3, "wo": 3, "w1": 3, "w2": 3, "w3": 3, "ffn": 3, "separ": [3, 47, 57, 164], "submodul": [3, 4, 190, 191], "feed_forward": 3, "parser": 3, "argumentpars": 3, "add_argu": 3, "torch_weight": 3, "output_fil": 3, "parse_arg": 3, "state": [3, 4, 175, 176, 190, 193, 194], "savez": 3, "k": [3, 191], "v": [3, 64, 191], "left": [3, 163, 170, 182, 183], "disk": 3, "text": 3, "format": [3, 95, 130, 131, 132], "oper": [3, 5, 33, 139, 145, 149, 190, 197, 198], "dictionari": [3, 175, 176, 178, 190, 191, 196], "represent": [3, 178, 180], "tree_unflatten": 3, "helper": 3, "weight_fil": 3, "incur": 3, "sever": [3, 62, 63, 131, 132], "unnecessari": [1, 3], "futur": 3, "pth": 3, "current": [3, 5, 6, 62, 63, 190], "around": 3, "m1": 3, "ultra": 3, "7b": 3, "me": 3, "ishmael": 3, "year": 3, "ago": 3, "never": 3, "long": 3, "info": 3, "247": 3, "press": 3, "enter": 3, "littl": 3, "monei": 3, "my": 3, "purs": 3, "greater": [3, 24, 91, 115], "consequ": 3, "walk": 3, "down": 3, "gower": 3, "street": 3, "afternoon": 3, "heavi": 3, "rain": 3, "saw": 3, "off": [3, 6], "man": 3, "rag": 3, "who": 3, "sat": 3, "upon": [3, 179], "hi": 3, "bundl": 3, "hard": 3, "wet": 3, "he": 3, "were": 3, "cry": 3, "watch": 3, "him": 3, "observ": 3, "numer": [3, 100, 102, 139, 164, 165, 168], "crowd": 3, "wa": [3, 176], "hurri": 3, "437": 3, "330": 3, "second": [3, 103, 153, 174], "spent": 3, "amount": 3, "39": 3, "ms": 3, "By": 3, "bigger": 3, "remain": [3, 153], "almost": 3, "nobodi": 3, "took": 3, "least": 3, "notic": 3, "distanc": 3, "had": 3, "doubt": 3, "minut": 3, "straight": 3, "slowli": 3, "rais": [3, 141], "ey": 3, "speak": 3, "resum": 3, "postur": 3, "stood": 3, "feel": 3, "pain": 3, "heart": 3, "smile": 3, "face": 3, "am": 3, "someon": 3, "three": 3, "quarter": 3, "hour": 3, "made": 3, "immedi": [3, 191], "repli": 3, "again": [3, 190], "hand": 3, "did": 3, "accustom": 3, "thu": [3, 190], "question": 3, "reason": 3, "tell": 3, "understand": 3, "579": 3, "690": 3, "num": [3, 124], "500": 3, "628": 3, "went": 3, "nervou": 3, "trembl": 3, "told": 3, "And": 3, "perhap": 3, "surpris": 3, "matter": [3, 190], "shall": 3, "anyhow": 3, "friend": 3, "ye": 3, "slight": 3, "kind": 3, "longer": [3, 64], "soon": 3, "unless": [3, 191], "unlik": [3, 13], "strang": 3, "amus": 3, "That": 3, "secret": 3, "disappoint": 3, "mine": 3, "cannot": 3, "happi": 3, "ask": 3, "Is": 3, "shop": 3, "bui": 3, "food": 3, "633": 3, "21": 3, "475": 3, "su": 3, "j": [3, 6, 174], "lu": 3, "pan": 3, "murtadha": 3, "wen": 3, "liu": 3, "2021": 3, "roform": 3, "enhanc": 3, "rotari": [3, 170], "arxiv": [3, 164, 165, 168, 170], "preprint": 3, "2104": [3, 170], "09864": [3, 170], "zhang": 3, "sennrich": 3, "2019": 3, "root": [3, 129, 142, 168], "advanc": 3, "inform": [3, 4, 163, 167], "system": 3, "shazeer": 3, "2020": 3, "glu": 3, "variant": 3, "2002": 3, "05202": 3, "classifi": 4, "mnist": 4, "As": [4, 148], "mlp": [4, 190, 193], "inherit": [4, 196], "standard": [4, 37, 56, 103, 119, 197], "idiom": 4, "input_dim": [4, 166], "hidden_dim": [4, 191, 193], "output_dim": [4, 166], "layer_s": 4, "idim": 4, "odim": 4, "maximum": [4, 22, 169, 182, 183, 184, 190, 191], "cross": 4, "entropi": 4, "sub": [4, 124], "commonli": [4, 191], "cross_entropi": 4, "accuraci": 4, "valid": [4, 64, 156, 178, 191, 196], "eval_fn": 4, "argmax": 4, "num_class": [4, 193], "batch_siz": [4, 193], "num_epoch": [4, 193], "learning_r": [4, 174, 177, 193], "train_imag": [4, 193], "train_label": [4, 193], "test_imag": 4, "test_label": 4, "shuffl": 4, "minibatch": 4, "batch_iter": [4, 193], "perm": 4, "permut": 4, "id": 4, "put": 4, "trainabl": [4, 173, 190, 191], "loss_and_grad_fn": [4, 193], "value_and_grad": [4, 190, 191, 193, 197], "epoch": 4, "test": [4, 6], "confus": 4, "decent": 4, "95": 4, "except": [5, 82, 83, 85, 86, 87, 164], "featur": [5, 62, 63, 164, 165, 166, 168, 170], "main": [5, 179, 190], "differ": [5, 146], "lazi": [5, 191, 197], "multi": [5, 160, 161], "cpu": 5, "gpu": 5, "strongli": [], "inspir": 5, "jax": [5, 194], "arrayfir": 5, "noteabl": 5, "unifi": 5, "live": 5, "guid": 5, "regress": 5, "layer": [5, 164, 165, 166, 171, 191], "perceptron": 5, "llm": 5, "infer": [5, 88], "fft": 5, "tree": [5, 73, 89, 153, 156, 175, 178, 179, 180], "develop": [5, 6], "document": [5, 47, 57], "17": 6, "g": [6, 177, 190, 198], "clang": 6, "cmake": 6, "24": 6, "clone": 6, "git": 6, "com": 6, "ml": 6, "explor": 6, "cd": 6, "brew": 6, "conda": 6, "global": [6, 123, 194], "env": 6, "cmake_build_parallel_level": 6, "edit": 6, "unittest": 6, "discov": 6, "mkdir": 6, "p": [6, 117, 174], "either": [6, 11, 47, 56, 57, 69, 70, 90, 91, 93, 94, 100, 103, 105, 108, 109, 146, 153, 171], "libmlx": 6, "preprocessor": 6, "metal_path": 6, "mlx_build_test": 6, "ON": 6, "mlx_build_exampl": 6, "mlx_build_benchmark": 6, "mlx_build_python_bind": 6, "devicetyp": 7, "attribut": [7, 8, 9, 26], "kwarg": [8, 26, 131, 132, 198], "union": [10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 28, 29, 30, 31, 32, 33, 34, 36, 38, 39, 40, 41, 42, 43, 45, 46, 47, 48, 50, 52, 53, 54, 55, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 69, 70, 71, 72, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 93, 94, 95, 96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 112, 113, 114, 115, 116, 117, 118, 119, 121, 122, 124, 125, 126, 127, 128, 129, 135, 136, 137, 138, 139, 140, 141, 142, 143, 144, 145, 146, 147, 148, 149, 150, 151, 152, 153, 154, 157, 158, 159, 161, 191], "wise": [10, 11, 16, 17, 18, 19, 20, 21, 65, 66, 69, 70, 71, 72, 74, 90, 91, 93, 94, 96, 97, 98, 99, 100, 101, 105, 108, 109, 110, 127, 129, 135, 136, 137, 138, 142, 143, 146, 150, 151, 172, 185], "absolut": [10, 13, 182, 183], "semant": [11, 60, 69, 70, 90, 91, 93, 94, 100, 103, 105, 108, 109, 146], "keepdim": [12, 14, 22, 23, 29, 30, 31, 32, 40, 41, 42, 43, 45, 55, 58, 102, 104, 106, 107, 116, 139, 147, 154], "reduct": [12, 14, 102, 104, 107, 116], "reduc": [12, 14, 22, 23, 102, 104, 106, 107, 116, 147, 154], "unspecifi": [12, 14, 15, 22, 23, 24, 25, 61, 88, 89, 102, 104, 106, 107, 112, 115, 116, 139, 140, 147, 148, 153, 154, 158, 198], "entir": [12, 14, 22, 23, 102, 104, 106, 107, 116, 147, 154], "singleton": [12, 14, 22, 23, 102, 103, 104, 106, 107, 116, 147, 154], "rtol": 13, "05": [13, 164, 165, 168], "atol": 13, "08": [13, 174], "approxim": [13, 163, 181, 182, 183], "comparison": [13, 70, 90, 91, 93, 94], "equal": [13, 24, 59, 91, 94, 115, 122, 141], "ab": [13, 153, 164, 165, 168, 170], "array_equ": 13, "rel": 13, "toler": 13, "boolean": [13, 59, 101, 187], "interv": [15, 122, 126], "increment": 15, "otherwis": [15, 178, 191], "int32": [15, 122, 187, 197], "convent": [15, 64], "lead": 15, "fraction": 15, "integr": [15, 148], "invers": [16, 17, 18, 19, 20, 21, 72, 79, 80, 81, 82, 83, 84], "cosin": [16, 17, 65, 66], "hyperbol": [17, 19, 21, 66, 138, 151], "sine": [18, 19, 137, 138], "minimum": [22, 23], "kth": [24, 115], "partit": 24, "order": [24, 115, 164, 171, 190, 191], "undefin": [24, 115], "sort": [24, 25, 115], "partiton": 24, "flatten": [24, 25, 115, 140, 148, 149, 178], "dimension": [26, 76, 77, 78, 79, 80, 81, 85, 86, 87, 160, 161, 162, 166], "overload": 26, "val": [26, 88], "tupl": [26, 47, 57, 61, 63, 73, 75, 92, 114, 128, 144, 153, 155, 161, 171, 178, 179, 180, 191, 196], "ndarrai": [26, 197], "properti": [27, 35, 44, 49, 51], "argument": [27, 47, 57, 73, 89, 130, 153, 179, 190, 194, 198], "elment": 51, "indices_or_sect": [52, 141], "nest": [56, 190, 191, 196], "correpsond": 56, "ddof": [58, 154], "equal_nan": 59, "nan": 59, "pad": [62, 63, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 160, 161], "dilat": [62, 63], "group": [62, 63, 164], "1d": [62, 64, 149], "convolut": [62, 63, 64, 160, 161], "channel": [62, 63, 160, 161], "c_in": [62, 63], "c_out": [62, 63], "convolv": [62, 63], "2d": 63, "spatial": [63, 164], "symmetr": 63, "discret": [64, 76, 77, 78, 79, 80, 81, 85, 86, 87, 162], "swap": 64, "conv": 64, "filter": [64, 160, 161, 191], "flip": 64, "signal": 64, "divis": 69, "quotient": 69, "mathrm": [71, 135], "frac": [71, 135, 164, 165, 168, 174], "pi": 71, "int_0": 71, "dx": 71, "erf": 72, "retain_graph": [73, 130], "node": [73, 156], "dict": [73, 95, 131, 191, 196], "leaf": [73, 178, 191], "preserv": [73, 128], "intend": 73, "control": [73, 194], "flow": [73, 145], "exponenti": 74, "insert": 75, "One": [76, 79, 85, 129], "fourier": [76, 77, 78, 79, 80, 81, 85, 86, 87], "truncat": [76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 125], "zero": [76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 159, 190], "dft": [76, 77, 78, 79, 80, 81, 85, 86, 87], "rfft": 82, "real": [82, 83, 84, 85, 86, 87], "rfft2": 83, "rfftn": 84, "silent": [85, 86, 87], "fun": [89, 92, 153, 155, 156], "cpp_function": [89, 153, 156], "variabl": [6, 89, 92, 153, 155, 156], "strict": [90, 93, 191], "binari": [95, 130, 131, 132], "npy": [95, 130], "natur": [96, 98], "logarithm": [96, 97, 98, 99], "log": [98, 100, 102], "plu": 98, "exp": [100, 102, 119, 139], "stabl": [100, 102, 139], "multipl": [6, 103, 109, 167], "prepend": 103, "remov": [103, 118, 144], "anoth": [103, 146, 157, 191], "negat": 110, "pad_width": 114, "constant_valu": 114, "edg": 114, "before_1": 114, "after_1": 114, "before_2": 114, "after_2": 114, "before_n": 114, "after_n": 114, "integ": [114, 117, 122, 141, 156, 162, 187], "before_i": 114, "after_i": 114, "extend": 114, "side": 114, "smaller": 115, "distribut": [117, 118, 119, 121, 125, 126], "prng": [117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 194], "num_sampl": 118, "unnorm": 118, "draw": 118, "uint32": [118, 187], "cdf": [119, 163, 181], "accord": [119, 157, 167], "seed": 120, "low": [122, 126], "high": [122, 126, 162, 190], "probabl": 122, "lower": [122, 125, 126], "upper": [122, 125, 126], "bound": [122, 125, 126, 163], "roadcast": 122, "domain": 125, "optino": 125, "uniformli": 126, "reciproc": 129, "arr": 130, "retain": 130, "dure": 130, "uncompress": 131, "my_path": 131, "tree_flatten": [131, 180], "transformerencod": 131, "128": [131, 190], "flat_param": 131, "keyword": [131, 132, 190, 194, 198], "compress": 132, "subarrai": 141, "being": [145, 190], "ident": 145, "prevent": 145, "unchang": [145, 170], "taken": 148, "prior": [148, 149], "equial": 148, "exclud": 149, "mse": 153, "param": [153, 190], "lvalu": 153, "dlvalu": 153, "dparam": 153, "lasso": 153, "l1": 153, "varianc": [154, 164], "divisor": 154, "cotang": 155, "in_ax": 156, "out_ax": 156, "prefix": [156, 178], "select": [157, 191], "in_channel": [160, 161], "out_channel": [160, 161], "kernel_s": [160, 161], "appli": [160, 161, 163, 164, 165, 166, 168, 169, 172, 175, 179, 181, 184, 185, 191], "nlc": 160, "learnabl": [160, 161, 171], "nhwc": 161, "height": 161, "width": 161, "num_embed": 162, "lookup": 162, "tabl": [162, 187], "typic": [162, 193], "usual": [162, 196], "vocabulari": 162, "approx": 163, "unit": [163, 169, 172, 181, 182, 183, 184, 185], "textrm": [163, 181], "phi": [163, 181], "geluapprox": 163, "sigma": [163, 172, 182, 183, 185], "60033": [163, 182], "0433603": [163, 182], "gelufast": 163, "773": [163, 183], "gelu_approx": [163, 181], "gelu_fast_approx": [163, 181], "regard": 163, "num_group": 164, "affin": [164, 165, 166], "pytorch_compat": 164, "var": [164, 165], "epsilon": [164, 165, 168, 174], "gamma": [164, 165, 168], "particular": 164, "split": 164, "preced": 164, "http": [164, 165, 168, 170], "org": [164, 165, 168, 170], "1803": 164, "08494": 164, "stabil": [164, 165, 168], "1607": 165, "06450": 165, "query_input_dim": 167, "key_input_dim": 167, "value_input_dim": 167, "value_dim": 167, "value_output_dim": 167, "dot": [167, 178, 191], "attent": [167, 191], "head": 167, "aggreg": 167, "lineari": 167, "bias": [167, 191], "inf": 167, "neg": 167, "attend": 167, "1910": 168, "07467": 168, "rectifi": [169, 184], "rotat": 170, "consecut": 170, "larger": 170, "slightli": 170, "callabl": [171, 173, 178, 179, 191], "plain": 171, "cdot": [172, 182, 183, 185], "fn": [173, 179, 197], "wrt": 173, "whose": 173, "9": 174, "999": 174, "paper": 174, "omit": 174, "estim": 174, "m_": 174, "beta_1": 174, "m_t": 174, "g_t": [174, 177], "v_": [174, 177], "beta_2": 174, "v_t": [174, 177], "w_": [174, 177], "w_t": [174, 177], "lambda": [174, 177, 179, 191], "kingma": 174, "ba": 174, "2015": 174, "stochast": [174, 177], "iclr": 174, "basi": 175, "optimizerst": 175, "recurs": [176, 190, 191], "defaultdict": 176, "miss": 176, "contrast": 176, "present": 176, "momentum": 177, "descent": 177, "mu": 177, "strength": 177, "is_leaf": 178, "notat": [178, 191], "arbitrari": [178, 191], "depth": 178, "hello": [178, 180], "charact": 178, "flat": [178, 180], "everi": 179, "superset": 179, "extra": 179, "closer": 179, "dict_kei": 179, "recreat": 180, "world": 180, "42": 180, "faster": 181, "gelu": [182, 183], "exact": [182, 183], "0003": 182, "015": 183, "show": [6, 187], "byte": 187, "bool_": 187, "uint8": 187, "unsign": 187, "uint16": 187, "16": [187, 191], "int8": 187, "sign": 187, "int16": 187, "int64": 187, "64": 187, "arm": 187, "arbitrarili": [190, 196, 197], "done": 190, "manual": 190, "explicitli": [190, 194], "solv": 190, "intuit": 190, "freez": [190, 191], "finetun": 190, "in_dim": [190, 191], "out_dim": [190, 191], "enumer": 190, "caus": 190, "local": 190, "scope": 190, "l2_loss": 190, "y_hat": 190, "trainable_paramet": [190, 191], "loss_and_grad": 190, "workhors": 190, "Its": 190, "frozen": [190, 191], "subset": [190, 191], "individu": 190, "action": 190, "preclud": 190, "pure": [190, 193], "pattern": 190, "achiev": 190, "other_input": 190, "necessari": 190, "wrap": 190, "subclass": 191, "concept": 191, "mymlp": 191, "in_proj": 191, "map_fn": 191, "filter_fn": 191, "valid_parameter_filt": 191, "apply_to_modul": 191, "apply_fn": 191, "children": 191, "descend": 191, "filter_and_map": 191, "is_leaf_fn": 191, "content": [6, 191], "found": 191, "whether": 191, "drop": 191, "idempot": 191, "ie": 191, "noop": 191, "unfreez": 191, "endswith": 191, "leaf_modul": 191, "load_weight": 191, "named_modul": 191, "save_weight": 191, "unfrozen": 191, "chang": 191, "tracer": 191, "partial": 191, "subsequ": 193, "implicit": 194, "fine": 194, "grain": 194, "manag": 194, "uniform": 194, "pseudo": 194, "altern": 194, "splittabl": 194, "threefri": 194, "counter": 194, "cycl": 196, "inspect": 197, "composit": 197, "sin": 197, "default_stream": 198, "default_devic": 198, "my_devic": 198, "brought": 5, "research": 5, "maco": 6, "13": 6, "recommend": 6, "14": 6, "sonoma": 6, "xcode": 6, "15": 6, "wish": 6, "environ": 6, "export": 6, "developer_dir": 6, "app": 6, "sdk": 6, "xcrun": 6, "macosx": 6}, "objects": {"mlx.core": [[7, 0, 1, "", "Device"], [8, 0, 1, "", "Dtype"], [9, 0, 1, "", "Stream"], [10, 2, 1, "", "abs"], [11, 2, 1, "", "add"], [12, 2, 1, "", "all"], [13, 2, 1, "", "allclose"], [14, 2, 1, "", "any"], [15, 2, 1, "", "arange"], [16, 2, 1, "", "arccos"], [17, 2, 1, "", "arccosh"], [18, 2, 1, "", "arcsin"], [19, 2, 1, "", "arcsinh"], [20, 2, 1, "", "arctan"], [21, 2, 1, "", "arctanh"], [22, 2, 1, "", "argmax"], [23, 2, 1, "", "argmin"], [24, 2, 1, "", "argpartition"], [25, 2, 1, "", "argsort"], [26, 0, 1, "", "array"], [59, 2, 1, "", "array_equal"], [60, 2, 1, "", "broadcast_to"], [61, 2, 1, "", "concatenate"], [62, 2, 1, "", "conv1d"], [63, 2, 1, "", "conv2d"], [64, 2, 1, "", "convolve"], [65, 2, 1, "", "cos"], [66, 2, 1, "", "cosh"], [67, 2, 1, "", "default_device"], [68, 2, 1, "", "default_stream"], [69, 2, 1, "", "divide"], [70, 2, 1, "", "equal"], [71, 2, 1, "", "erf"], [72, 2, 1, "", "erfinv"], [73, 2, 1, "", "eval"], [74, 2, 1, "", "exp"], [75, 2, 1, "", "expand_dims"], [88, 2, 1, "", "full"], [89, 2, 1, "", "grad"], [90, 2, 1, "", "greater"], [91, 2, 1, "", "greater_equal"], [92, 2, 1, "", "jvp"], [93, 2, 1, "", "less"], [94, 2, 1, "", "less_equal"], [95, 2, 1, "", "load"], [96, 2, 1, "", "log"], [97, 2, 1, "", "log10"], [98, 2, 1, "", "log1p"], [99, 2, 1, "", "log2"], [100, 2, 1, "", "logaddexp"], [101, 2, 1, "", "logical_not"], [102, 2, 1, "", "logsumexp"], [103, 2, 1, "", "matmul"], [104, 2, 1, "", "max"], [105, 2, 1, "", "maximum"], [106, 2, 1, "", "mean"], [107, 2, 1, "", "min"], [108, 2, 1, "", "minimum"], [109, 2, 1, "", "multiply"], [110, 2, 1, "", "negative"], [111, 2, 1, "", "new_stream"], [112, 2, 1, "", "ones"], [113, 2, 1, "", "ones_like"], [114, 2, 1, "", "pad"], [115, 2, 1, "", "partition"], [116, 2, 1, "", "prod"], [127, 2, 1, "", "reciprocal"], [128, 2, 1, "", "reshape"], [129, 2, 1, "", "rsqrt"], [130, 2, 1, "", "save"], [131, 2, 1, "", "savez"], [132, 2, 1, "", "savez_compressed"], [133, 2, 1, "", "set_default_device"], [134, 2, 1, "", "set_default_stream"], [135, 2, 1, "", "sigmoid"], [136, 2, 1, "", "sign"], [137, 2, 1, "", "sin"], [138, 2, 1, "", "sinh"], [139, 2, 1, "", "softmax"], [140, 2, 1, "", "sort"], [141, 2, 1, "", "split"], [142, 2, 1, "", "sqrt"], [143, 2, 1, "", "square"], [144, 2, 1, "", "squeeze"], [145, 2, 1, "", "stop_gradient"], [146, 2, 1, "", "subtract"], [147, 2, 1, "", "sum"], [148, 2, 1, "", "take"], [149, 2, 1, "", "take_along_axis"], [150, 2, 1, "", "tan"], [151, 2, 1, "", "tanh"], [152, 2, 1, "", "transpose"], [153, 2, 1, "", "value_and_grad"], [154, 2, 1, "", "var"], [155, 2, 1, "", "vjp"], [156, 2, 1, "", "vmap"], [157, 2, 1, "", "where"], [158, 2, 1, "", "zeros"], [159, 2, 1, "", "zeros_like"]], "mlx.core.Device": [[7, 1, 1, "", "__init__"]], "mlx.core.Dtype": [[8, 1, 1, "", "__init__"]], "mlx.core.Stream": [[9, 1, 1, "", "__init__"]], "mlx.core.array": [[27, 3, 1, "", "T"], [26, 1, 1, "", "__init__"], [28, 1, 1, "", "abs"], [29, 1, 1, "", "all"], [30, 1, 1, "", "any"], [31, 1, 1, "", "argmax"], [32, 1, 1, "", "argmin"], [33, 1, 1, "", "astype"], [34, 1, 1, "", "cos"], [35, 3, 1, "", "dtype"], [36, 1, 1, "", "exp"], [37, 1, 1, "", "item"], [38, 1, 1, "", "log"], [39, 1, 1, "", "log1p"], [40, 1, 1, "", "logsumexp"], [41, 1, 1, "", "max"], [42, 1, 1, "", "mean"], [43, 1, 1, "", "min"], [44, 3, 1, "", "ndim"], [45, 1, 1, "", "prod"], [46, 1, 1, "", "reciprocal"], [47, 1, 1, "", "reshape"], [48, 1, 1, "", "rsqrt"], [49, 3, 1, "", "shape"], [50, 1, 1, "", "sin"], [51, 3, 1, "", "size"], [52, 1, 1, "", "split"], [53, 1, 1, "", "sqrt"], [54, 1, 1, "", "square"], [55, 1, 1, "", "sum"], [56, 1, 1, "", "tolist"], [57, 1, 1, "", "transpose"], [58, 1, 1, "", "var"]], "mlx.core.fft": [[76, 2, 1, "", "fft"], [77, 2, 1, "", "fft2"], [78, 2, 1, "", "fftn"], [79, 2, 1, "", "ifft"], [80, 2, 1, "", "ifft2"], [81, 2, 1, "", "ifftn"], [82, 2, 1, "", "irfft"], [83, 2, 1, "", "irfft2"], [84, 2, 1, "", "irfftn"], [85, 2, 1, "", "rfft"], [86, 2, 1, "", "rfft2"], [87, 2, 1, "", "rfftn"]], "mlx.core.random": [[117, 2, 1, "", "bernoulli"], [118, 2, 1, "", "categorical"], [119, 2, 1, "", "gumbel"], [120, 2, 1, "", "key"], [121, 2, 1, "", "normal"], [122, 2, 1, "", "randint"], [123, 2, 1, "", "seed"], [124, 2, 1, "", "split"], [125, 2, 1, "", "truncated_normal"], [126, 2, 1, "", "uniform"]], "mlx.nn": [[160, 0, 1, "", "Conv1d"], [161, 0, 1, "", "Conv2d"], [162, 0, 1, "", "Embedding"], [163, 0, 1, "", "GELU"], [164, 0, 1, "", "GroupNorm"], [165, 0, 1, "", "LayerNorm"], [166, 0, 1, "", "Linear"], [191, 0, 1, "", "Module"], [167, 0, 1, "", "MultiHeadAttention"], [168, 0, 1, "", "RMSNorm"], [169, 0, 1, "", "ReLU"], [170, 0, 1, "", "RoPE"], [171, 0, 1, "", "Sequential"], [172, 0, 1, "", "SiLU"], [181, 0, 1, "", "gelu"], [182, 0, 1, "", "gelu_approx"], [183, 0, 1, "", "gelu_fast_approx"], [184, 0, 1, "", "relu"], [185, 0, 1, "", "silu"], [173, 2, 1, "", "value_and_grad"]], "mlx.nn.Module": [[191, 1, 1, "", "apply"], [191, 1, 1, "", "apply_to_modules"], [191, 1, 1, "", "children"], [191, 1, 1, "", "filter_and_map"], [191, 1, 1, "", "freeze"], [191, 1, 1, "", "leaf_modules"], [191, 1, 1, "", "load_weights"], [191, 1, 1, "", "modules"], [191, 1, 1, "", "named_modules"], [191, 1, 1, "", "parameters"], [191, 1, 1, "", "save_weights"], [191, 1, 1, "", "trainable_parameters"], [191, 1, 1, "", "unfreeze"], [191, 1, 1, "", "update"]], "mlx.optimizers": [[174, 0, 1, "", "Adam"], [175, 0, 1, "", "Optimizer"], [176, 0, 1, "", "OptimizerState"], [177, 0, 1, "", "SGD"]], "mlx.optimizers.Optimizer": [[175, 4, 1, "", "state"]], "mlx.utils": [[178, 2, 1, "", "tree_flatten"], [179, 2, 1, "", "tree_map"], [180, 2, 1, "", "tree_unflatten"]]}, "objtypes": {"0": "py:class", "1": "py:method", "2": "py:function", "3": "py:property", "4": "py:attribute"}, "objnames": {"0": ["py", "class", "Python class"], "1": ["py", "method", "Python method"], "2": ["py", "function", "Python function"], "3": ["py", "property", "Python property"], "4": ["py", "attribute", "Python attribute"]}, "titleterms": {"oper": [0, 1, 192], "develop": 1, "document": 1, "introduc": 1, "exampl": [1, 5], "primit": 1, "us": [1, 198], "implement": [1, 3], "cpu": 1, "backend": 1, "gpu": 1, "transform": [1, 195, 197], "build": [1, 6], "bind": 1, "python": [1, 5, 6], "cmake": 1, "setuptool": 1, "usag": [1, 5], "result": 1, "script": [1, 3], "download": [1, 3], "code": [1, 3], "linear": [2, 166], "regress": 2, "llm": 3, "infer": 3, "model": 3, "attent": 3, "layer": [3, 4, 190], "encod": 3, "full": [3, 88], "gener": 3, "put": 3, "all": [3, 12, 29], "togeth": 3, "convert": 3, "weight": 3, "load": [3, 95], "benchmark": 3, "multi": 4, "perceptron": 4, "mlx": [5, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 127, 128, 129, 130, 131, 132, 133, 134, 135, 136, 137, 138, 139, 140, 141, 142, 143, 144, 145, 146, 147, 148, 149, 150, 151, 152, 153, 154, 155, 156, 157, 158, 159, 160, 161, 162, 163, 164, 165, 166, 167, 168, 169, 170, 171, 172, 173, 174, 175, 176, 177, 178, 179, 180, 181, 182, 183, 184, 185, 191], "instal": [5, 6], "api": [5, 6], "refer": 5, "c": [5, 6], "further": 5, "read": 5, "from": 6, "pypi": 6, "sourc": 6, "requir": 6, "option": 6, "core": [7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 127, 128, 129, 130, 131, 132, 133, 134, 135, 136, 137, 138, 139, 140, 141, 142, 143, 144, 145, 146, 147, 148, 149, 150, 151, 152, 153, 154, 155, 156, 157, 158, 159], "devic": [7, 188, 197], "dtype": [8, 35], "stream": [9, 188, 197, 198], "ab": [10, 28], "add": 11, "allclos": 13, "ani": [14, 30], "arang": 15, "arcco": 16, "arccosh": 17, "arcsin": 18, "arcsinh": 19, "arctan": 20, "arctanh": 21, "argmax": [22, 31], "argmin": [23, 32], "argpartit": 24, "argsort": 25, "arrai": [26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 186], "t": 27, "astyp": 33, "co": [34, 65], "exp": [36, 74], "item": 37, "log": [38, 96], "log1p": [39, 98], "logsumexp": [40, 102], "max": [41, 104], "mean": [42, 106], "min": [43, 107], "ndim": 44, "prod": [45, 116], "reciproc": [46, 127], "reshap": [47, 128], "rsqrt": [48, 129], "shape": 49, "sin": [50, 137], "size": 51, "split": [52, 124, 141], "sqrt": [53, 142], "squar": [54, 143], "sum": [55, 147], "tolist": 56, "transpos": [57, 152], "var": [58, 154], "array_equ": 59, "broadcast_to": 60, "concaten": 61, "conv1d": [62, 160], "conv2d": [63, 161], "convolv": 64, "cosh": 66, "default_devic": 67, "default_stream": 68, "divid": 69, "equal": 70, "erf": 71, "erfinv": 72, "eval": 73, "expand_dim": 75, "fft": [76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 189], "fft2": 77, "fftn": 78, "ifft": 79, "ifft2": 80, "ifftn": 81, "irfft": 82, "irfft2": 83, "irfftn": 84, "rfft": 85, "rfft2": 86, "rfftn": 87, "grad": [89, 190], "greater": 90, "greater_equ": 91, "jvp": 92, "less": 93, "less_equ": 94, "log10": 97, "log2": 99, "logaddexp": 100, "logical_not": 101, "matmul": 103, "maximum": 105, "minimum": 108, "multipli": 109, "neg": 110, "new_stream": 111, "ones": 112, "ones_lik": 113, "pad": 114, "partit": 115, "random": [117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 194], "bernoulli": 117, "categor": 118, "gumbel": 119, "kei": 120, "normal": 121, "randint": 122, "seed": 123, "truncated_norm": 125, "uniform": 126, "save": 130, "savez": 131, "savez_compress": 132, "set_default_devic": 133, "set_default_stream": 134, "sigmoid": 135, "sign": 136, "sinh": 138, "softmax": 139, "sort": 140, "squeez": 144, "stop_gradi": 145, "subtract": 146, "take": 148, "take_along_axi": 149, "tan": 150, "tanh": 151, "value_and_grad": [153, 173], "vjp": 155, "vmap": 156, "where": 157, "zero": 158, "zeros_lik": 159, "nn": [160, 161, 162, 163, 164, 165, 166, 167, 168, 169, 170, 171, 172, 173, 181, 182, 183, 184, 185, 191], "embed": 162, "gelu": [163, 181], "groupnorm": 164, "layernorm": 165, "multiheadattent": 167, "rmsnorm": 168, "relu": [169, 184], "rope": 170, "sequenti": 171, "silu": [172, 185], "optim": [174, 175, 176, 177, 193], "adam": 174, "optimizerst": 176, "sgd": 177, "util": [178, 179, 180, 196], "tree_flatten": 178, "tree_map": 179, "tree_unflatten": 180, "gelu_approx": 182, "gelu_fast_approx": 183, "data": 187, "type": 187, "support": 187, "neural": 190, "network": 190, "quick": [190, 197], "start": [190, 197], "The": 190, "modul": [190, 191], "class": 190, "paramet": 190, "updat": 190, "valu": 190, "tree": 196, "guid": 197, "basic": 197, "function": 197, "graph": 197, "specifi": 198}, "envversion": {"sphinx.domains.c": 2, "sphinx.domains.changeset": 1, "sphinx.domains.citation": 1, "sphinx.domains.cpp": 6, "sphinx.domains.index": 1, "sphinx.domains.javascript": 2, "sphinx.domains.math": 2, "sphinx.domains.python": 3, "sphinx.domains.rst": 2, "sphinx.domains.std": 2, "sphinx.ext.intersphinx": 1, "sphinx": 56}})
\ No newline at end of file
+Search.setIndex({"docnames": ["cpp/ops", "dev/extensions", "examples/linear_regression", "examples/llama-inference", "examples/mlp", "index", "install", "python/_autosummary/mlx.core.Device", "python/_autosummary/mlx.core.Dtype", "python/_autosummary/mlx.core.Stream", "python/_autosummary/mlx.core.abs", "python/_autosummary/mlx.core.add", "python/_autosummary/mlx.core.all", "python/_autosummary/mlx.core.allclose", "python/_autosummary/mlx.core.any", "python/_autosummary/mlx.core.arange", "python/_autosummary/mlx.core.arccos", "python/_autosummary/mlx.core.arccosh", "python/_autosummary/mlx.core.arcsin", "python/_autosummary/mlx.core.arcsinh", "python/_autosummary/mlx.core.arctan", "python/_autosummary/mlx.core.arctanh", "python/_autosummary/mlx.core.argmax", "python/_autosummary/mlx.core.argmin", "python/_autosummary/mlx.core.argpartition", "python/_autosummary/mlx.core.argsort", "python/_autosummary/mlx.core.array", "python/_autosummary/mlx.core.array.T", "python/_autosummary/mlx.core.array.abs", "python/_autosummary/mlx.core.array.all", "python/_autosummary/mlx.core.array.any", "python/_autosummary/mlx.core.array.argmax", "python/_autosummary/mlx.core.array.argmin", "python/_autosummary/mlx.core.array.astype", "python/_autosummary/mlx.core.array.cos", "python/_autosummary/mlx.core.array.dtype", "python/_autosummary/mlx.core.array.exp", "python/_autosummary/mlx.core.array.item", "python/_autosummary/mlx.core.array.log", "python/_autosummary/mlx.core.array.log1p", "python/_autosummary/mlx.core.array.logsumexp", "python/_autosummary/mlx.core.array.max", "python/_autosummary/mlx.core.array.mean", "python/_autosummary/mlx.core.array.min", "python/_autosummary/mlx.core.array.ndim", "python/_autosummary/mlx.core.array.prod", "python/_autosummary/mlx.core.array.reciprocal", "python/_autosummary/mlx.core.array.reshape", "python/_autosummary/mlx.core.array.rsqrt", "python/_autosummary/mlx.core.array.shape", "python/_autosummary/mlx.core.array.sin", "python/_autosummary/mlx.core.array.size", "python/_autosummary/mlx.core.array.split", "python/_autosummary/mlx.core.array.sqrt", "python/_autosummary/mlx.core.array.square", "python/_autosummary/mlx.core.array.sum", "python/_autosummary/mlx.core.array.tolist", "python/_autosummary/mlx.core.array.transpose", "python/_autosummary/mlx.core.array.var", "python/_autosummary/mlx.core.array_equal", "python/_autosummary/mlx.core.broadcast_to", "python/_autosummary/mlx.core.concatenate", "python/_autosummary/mlx.core.conv1d", "python/_autosummary/mlx.core.conv2d", "python/_autosummary/mlx.core.convolve", "python/_autosummary/mlx.core.cos", "python/_autosummary/mlx.core.cosh", "python/_autosummary/mlx.core.default_device", "python/_autosummary/mlx.core.default_stream", "python/_autosummary/mlx.core.divide", "python/_autosummary/mlx.core.equal", "python/_autosummary/mlx.core.erf", "python/_autosummary/mlx.core.erfinv", "python/_autosummary/mlx.core.eval", "python/_autosummary/mlx.core.exp", "python/_autosummary/mlx.core.expand_dims", "python/_autosummary/mlx.core.fft.fft", "python/_autosummary/mlx.core.fft.fft2", "python/_autosummary/mlx.core.fft.fftn", "python/_autosummary/mlx.core.fft.ifft", "python/_autosummary/mlx.core.fft.ifft2", "python/_autosummary/mlx.core.fft.ifftn", "python/_autosummary/mlx.core.fft.irfft", "python/_autosummary/mlx.core.fft.irfft2", "python/_autosummary/mlx.core.fft.irfftn", "python/_autosummary/mlx.core.fft.rfft", "python/_autosummary/mlx.core.fft.rfft2", "python/_autosummary/mlx.core.fft.rfftn", "python/_autosummary/mlx.core.full", "python/_autosummary/mlx.core.grad", "python/_autosummary/mlx.core.greater", "python/_autosummary/mlx.core.greater_equal", "python/_autosummary/mlx.core.jvp", "python/_autosummary/mlx.core.less", "python/_autosummary/mlx.core.less_equal", "python/_autosummary/mlx.core.load", "python/_autosummary/mlx.core.log", "python/_autosummary/mlx.core.log10", "python/_autosummary/mlx.core.log1p", "python/_autosummary/mlx.core.log2", "python/_autosummary/mlx.core.logaddexp", "python/_autosummary/mlx.core.logical_not", "python/_autosummary/mlx.core.logsumexp", "python/_autosummary/mlx.core.matmul", "python/_autosummary/mlx.core.max", "python/_autosummary/mlx.core.maximum", "python/_autosummary/mlx.core.mean", "python/_autosummary/mlx.core.min", "python/_autosummary/mlx.core.minimum", "python/_autosummary/mlx.core.multiply", "python/_autosummary/mlx.core.negative", "python/_autosummary/mlx.core.new_stream", "python/_autosummary/mlx.core.ones", "python/_autosummary/mlx.core.ones_like", "python/_autosummary/mlx.core.pad", "python/_autosummary/mlx.core.partition", "python/_autosummary/mlx.core.prod", "python/_autosummary/mlx.core.random.bernoulli", "python/_autosummary/mlx.core.random.categorical", "python/_autosummary/mlx.core.random.gumbel", "python/_autosummary/mlx.core.random.key", "python/_autosummary/mlx.core.random.normal", "python/_autosummary/mlx.core.random.randint", "python/_autosummary/mlx.core.random.seed", "python/_autosummary/mlx.core.random.split", "python/_autosummary/mlx.core.random.truncated_normal", "python/_autosummary/mlx.core.random.uniform", "python/_autosummary/mlx.core.reciprocal", "python/_autosummary/mlx.core.reshape", "python/_autosummary/mlx.core.rsqrt", "python/_autosummary/mlx.core.save", "python/_autosummary/mlx.core.savez", "python/_autosummary/mlx.core.savez_compressed", "python/_autosummary/mlx.core.set_default_device", "python/_autosummary/mlx.core.set_default_stream", "python/_autosummary/mlx.core.sigmoid", "python/_autosummary/mlx.core.sign", "python/_autosummary/mlx.core.sin", "python/_autosummary/mlx.core.sinh", "python/_autosummary/mlx.core.softmax", "python/_autosummary/mlx.core.sort", "python/_autosummary/mlx.core.split", "python/_autosummary/mlx.core.sqrt", "python/_autosummary/mlx.core.square", "python/_autosummary/mlx.core.squeeze", "python/_autosummary/mlx.core.stop_gradient", "python/_autosummary/mlx.core.subtract", "python/_autosummary/mlx.core.sum", "python/_autosummary/mlx.core.take", "python/_autosummary/mlx.core.take_along_axis", "python/_autosummary/mlx.core.tan", "python/_autosummary/mlx.core.tanh", "python/_autosummary/mlx.core.transpose", "python/_autosummary/mlx.core.value_and_grad", "python/_autosummary/mlx.core.var", "python/_autosummary/mlx.core.vjp", "python/_autosummary/mlx.core.vmap", "python/_autosummary/mlx.core.where", "python/_autosummary/mlx.core.zeros", "python/_autosummary/mlx.core.zeros_like", "python/_autosummary/mlx.nn.Conv1d", "python/_autosummary/mlx.nn.Conv2d", "python/_autosummary/mlx.nn.Embedding", "python/_autosummary/mlx.nn.GELU", "python/_autosummary/mlx.nn.GroupNorm", "python/_autosummary/mlx.nn.LayerNorm", "python/_autosummary/mlx.nn.Linear", "python/_autosummary/mlx.nn.MultiHeadAttention", "python/_autosummary/mlx.nn.RMSNorm", "python/_autosummary/mlx.nn.ReLU", "python/_autosummary/mlx.nn.RoPE", "python/_autosummary/mlx.nn.Sequential", "python/_autosummary/mlx.nn.SiLU", "python/_autosummary/mlx.nn.value_and_grad", "python/_autosummary/mlx.optimizers.Adam", "python/_autosummary/mlx.optimizers.Optimizer", "python/_autosummary/mlx.optimizers.OptimizerState", "python/_autosummary/mlx.optimizers.SGD", "python/_autosummary/mlx.utils.tree_flatten", "python/_autosummary/mlx.utils.tree_map", "python/_autosummary/mlx.utils.tree_unflatten", "python/_autosummary_functions/mlx.nn.gelu", "python/_autosummary_functions/mlx.nn.gelu_approx", "python/_autosummary_functions/mlx.nn.gelu_fast_approx", "python/_autosummary_functions/mlx.nn.relu", "python/_autosummary_functions/mlx.nn.silu", "python/array", "python/data_types", "python/devices_and_streams", "python/fft", "python/nn", "python/nn/module", "python/ops", "python/optimizers", "python/random", "python/transforms", "python/tree_utils", "quick_start", "using_streams"], "filenames": ["cpp/ops.rst", "dev/extensions.rst", "examples/linear_regression.rst", "examples/llama-inference.rst", "examples/mlp.rst", "index.rst", "install.rst", "python/_autosummary/mlx.core.Device.rst", "python/_autosummary/mlx.core.Dtype.rst", "python/_autosummary/mlx.core.Stream.rst", "python/_autosummary/mlx.core.abs.rst", "python/_autosummary/mlx.core.add.rst", "python/_autosummary/mlx.core.all.rst", "python/_autosummary/mlx.core.allclose.rst", "python/_autosummary/mlx.core.any.rst", "python/_autosummary/mlx.core.arange.rst", "python/_autosummary/mlx.core.arccos.rst", "python/_autosummary/mlx.core.arccosh.rst", "python/_autosummary/mlx.core.arcsin.rst", "python/_autosummary/mlx.core.arcsinh.rst", "python/_autosummary/mlx.core.arctan.rst", "python/_autosummary/mlx.core.arctanh.rst", "python/_autosummary/mlx.core.argmax.rst", "python/_autosummary/mlx.core.argmin.rst", "python/_autosummary/mlx.core.argpartition.rst", "python/_autosummary/mlx.core.argsort.rst", "python/_autosummary/mlx.core.array.rst", "python/_autosummary/mlx.core.array.T.rst", "python/_autosummary/mlx.core.array.abs.rst", "python/_autosummary/mlx.core.array.all.rst", "python/_autosummary/mlx.core.array.any.rst", "python/_autosummary/mlx.core.array.argmax.rst", "python/_autosummary/mlx.core.array.argmin.rst", "python/_autosummary/mlx.core.array.astype.rst", "python/_autosummary/mlx.core.array.cos.rst", "python/_autosummary/mlx.core.array.dtype.rst", "python/_autosummary/mlx.core.array.exp.rst", "python/_autosummary/mlx.core.array.item.rst", "python/_autosummary/mlx.core.array.log.rst", "python/_autosummary/mlx.core.array.log1p.rst", "python/_autosummary/mlx.core.array.logsumexp.rst", "python/_autosummary/mlx.core.array.max.rst", "python/_autosummary/mlx.core.array.mean.rst", "python/_autosummary/mlx.core.array.min.rst", "python/_autosummary/mlx.core.array.ndim.rst", "python/_autosummary/mlx.core.array.prod.rst", "python/_autosummary/mlx.core.array.reciprocal.rst", "python/_autosummary/mlx.core.array.reshape.rst", "python/_autosummary/mlx.core.array.rsqrt.rst", "python/_autosummary/mlx.core.array.shape.rst", "python/_autosummary/mlx.core.array.sin.rst", "python/_autosummary/mlx.core.array.size.rst", "python/_autosummary/mlx.core.array.split.rst", "python/_autosummary/mlx.core.array.sqrt.rst", "python/_autosummary/mlx.core.array.square.rst", "python/_autosummary/mlx.core.array.sum.rst", "python/_autosummary/mlx.core.array.tolist.rst", "python/_autosummary/mlx.core.array.transpose.rst", "python/_autosummary/mlx.core.array.var.rst", "python/_autosummary/mlx.core.array_equal.rst", "python/_autosummary/mlx.core.broadcast_to.rst", "python/_autosummary/mlx.core.concatenate.rst", "python/_autosummary/mlx.core.conv1d.rst", "python/_autosummary/mlx.core.conv2d.rst", "python/_autosummary/mlx.core.convolve.rst", "python/_autosummary/mlx.core.cos.rst", "python/_autosummary/mlx.core.cosh.rst", "python/_autosummary/mlx.core.default_device.rst", "python/_autosummary/mlx.core.default_stream.rst", "python/_autosummary/mlx.core.divide.rst", "python/_autosummary/mlx.core.equal.rst", "python/_autosummary/mlx.core.erf.rst", "python/_autosummary/mlx.core.erfinv.rst", "python/_autosummary/mlx.core.eval.rst", "python/_autosummary/mlx.core.exp.rst", "python/_autosummary/mlx.core.expand_dims.rst", "python/_autosummary/mlx.core.fft.fft.rst", "python/_autosummary/mlx.core.fft.fft2.rst", "python/_autosummary/mlx.core.fft.fftn.rst", "python/_autosummary/mlx.core.fft.ifft.rst", "python/_autosummary/mlx.core.fft.ifft2.rst", "python/_autosummary/mlx.core.fft.ifftn.rst", "python/_autosummary/mlx.core.fft.irfft.rst", "python/_autosummary/mlx.core.fft.irfft2.rst", "python/_autosummary/mlx.core.fft.irfftn.rst", "python/_autosummary/mlx.core.fft.rfft.rst", "python/_autosummary/mlx.core.fft.rfft2.rst", "python/_autosummary/mlx.core.fft.rfftn.rst", "python/_autosummary/mlx.core.full.rst", "python/_autosummary/mlx.core.grad.rst", "python/_autosummary/mlx.core.greater.rst", "python/_autosummary/mlx.core.greater_equal.rst", "python/_autosummary/mlx.core.jvp.rst", "python/_autosummary/mlx.core.less.rst", "python/_autosummary/mlx.core.less_equal.rst", "python/_autosummary/mlx.core.load.rst", "python/_autosummary/mlx.core.log.rst", "python/_autosummary/mlx.core.log10.rst", "python/_autosummary/mlx.core.log1p.rst", "python/_autosummary/mlx.core.log2.rst", "python/_autosummary/mlx.core.logaddexp.rst", "python/_autosummary/mlx.core.logical_not.rst", "python/_autosummary/mlx.core.logsumexp.rst", "python/_autosummary/mlx.core.matmul.rst", "python/_autosummary/mlx.core.max.rst", "python/_autosummary/mlx.core.maximum.rst", "python/_autosummary/mlx.core.mean.rst", "python/_autosummary/mlx.core.min.rst", "python/_autosummary/mlx.core.minimum.rst", "python/_autosummary/mlx.core.multiply.rst", "python/_autosummary/mlx.core.negative.rst", "python/_autosummary/mlx.core.new_stream.rst", "python/_autosummary/mlx.core.ones.rst", "python/_autosummary/mlx.core.ones_like.rst", "python/_autosummary/mlx.core.pad.rst", "python/_autosummary/mlx.core.partition.rst", "python/_autosummary/mlx.core.prod.rst", "python/_autosummary/mlx.core.random.bernoulli.rst", "python/_autosummary/mlx.core.random.categorical.rst", "python/_autosummary/mlx.core.random.gumbel.rst", "python/_autosummary/mlx.core.random.key.rst", "python/_autosummary/mlx.core.random.normal.rst", "python/_autosummary/mlx.core.random.randint.rst", "python/_autosummary/mlx.core.random.seed.rst", "python/_autosummary/mlx.core.random.split.rst", "python/_autosummary/mlx.core.random.truncated_normal.rst", "python/_autosummary/mlx.core.random.uniform.rst", "python/_autosummary/mlx.core.reciprocal.rst", "python/_autosummary/mlx.core.reshape.rst", "python/_autosummary/mlx.core.rsqrt.rst", "python/_autosummary/mlx.core.save.rst", "python/_autosummary/mlx.core.savez.rst", "python/_autosummary/mlx.core.savez_compressed.rst", "python/_autosummary/mlx.core.set_default_device.rst", "python/_autosummary/mlx.core.set_default_stream.rst", "python/_autosummary/mlx.core.sigmoid.rst", "python/_autosummary/mlx.core.sign.rst", "python/_autosummary/mlx.core.sin.rst", "python/_autosummary/mlx.core.sinh.rst", "python/_autosummary/mlx.core.softmax.rst", "python/_autosummary/mlx.core.sort.rst", "python/_autosummary/mlx.core.split.rst", "python/_autosummary/mlx.core.sqrt.rst", "python/_autosummary/mlx.core.square.rst", "python/_autosummary/mlx.core.squeeze.rst", "python/_autosummary/mlx.core.stop_gradient.rst", "python/_autosummary/mlx.core.subtract.rst", "python/_autosummary/mlx.core.sum.rst", "python/_autosummary/mlx.core.take.rst", "python/_autosummary/mlx.core.take_along_axis.rst", "python/_autosummary/mlx.core.tan.rst", "python/_autosummary/mlx.core.tanh.rst", "python/_autosummary/mlx.core.transpose.rst", "python/_autosummary/mlx.core.value_and_grad.rst", "python/_autosummary/mlx.core.var.rst", "python/_autosummary/mlx.core.vjp.rst", "python/_autosummary/mlx.core.vmap.rst", "python/_autosummary/mlx.core.where.rst", "python/_autosummary/mlx.core.zeros.rst", "python/_autosummary/mlx.core.zeros_like.rst", "python/_autosummary/mlx.nn.Conv1d.rst", "python/_autosummary/mlx.nn.Conv2d.rst", "python/_autosummary/mlx.nn.Embedding.rst", "python/_autosummary/mlx.nn.GELU.rst", "python/_autosummary/mlx.nn.GroupNorm.rst", "python/_autosummary/mlx.nn.LayerNorm.rst", "python/_autosummary/mlx.nn.Linear.rst", "python/_autosummary/mlx.nn.MultiHeadAttention.rst", "python/_autosummary/mlx.nn.RMSNorm.rst", "python/_autosummary/mlx.nn.ReLU.rst", "python/_autosummary/mlx.nn.RoPE.rst", "python/_autosummary/mlx.nn.Sequential.rst", "python/_autosummary/mlx.nn.SiLU.rst", "python/_autosummary/mlx.nn.value_and_grad.rst", "python/_autosummary/mlx.optimizers.Adam.rst", "python/_autosummary/mlx.optimizers.Optimizer.rst", "python/_autosummary/mlx.optimizers.OptimizerState.rst", "python/_autosummary/mlx.optimizers.SGD.rst", "python/_autosummary/mlx.utils.tree_flatten.rst", "python/_autosummary/mlx.utils.tree_map.rst", "python/_autosummary/mlx.utils.tree_unflatten.rst", "python/_autosummary_functions/mlx.nn.gelu.rst", "python/_autosummary_functions/mlx.nn.gelu_approx.rst", "python/_autosummary_functions/mlx.nn.gelu_fast_approx.rst", "python/_autosummary_functions/mlx.nn.relu.rst", "python/_autosummary_functions/mlx.nn.silu.rst", "python/array.rst", "python/data_types.rst", "python/devices_and_streams.rst", "python/fft.rst", "python/nn.rst", "python/nn/module.rst", "python/ops.rst", "python/optimizers.rst", "python/random.rst", "python/transforms.rst", "python/tree_utils.rst", "quick_start.rst", "using_streams.rst"], "titles": ["Operations", "Developer Documentation", "Linear Regression", "LLM inference", "Multi-Layer Perceptron", "MLX", "Build and Install", "mlx.core.Device", "mlx.core.Dtype", "mlx.core.Stream", "mlx.core.abs", "mlx.core.add", "mlx.core.all", "mlx.core.allclose", "mlx.core.any", "mlx.core.arange", "mlx.core.arccos", "mlx.core.arccosh", "mlx.core.arcsin", "mlx.core.arcsinh", "mlx.core.arctan", "mlx.core.arctanh", "mlx.core.argmax", "mlx.core.argmin", "mlx.core.argpartition", "mlx.core.argsort", "mlx.core.array", "mlx.core.array.T", "mlx.core.array.abs", "mlx.core.array.all", "mlx.core.array.any", "mlx.core.array.argmax", "mlx.core.array.argmin", "mlx.core.array.astype", "mlx.core.array.cos", "mlx.core.array.dtype", "mlx.core.array.exp", "mlx.core.array.item", "mlx.core.array.log", "mlx.core.array.log1p", "mlx.core.array.logsumexp", "mlx.core.array.max", "mlx.core.array.mean", "mlx.core.array.min", "mlx.core.array.ndim", "mlx.core.array.prod", "mlx.core.array.reciprocal", "mlx.core.array.reshape", "mlx.core.array.rsqrt", "mlx.core.array.shape", "mlx.core.array.sin", "mlx.core.array.size", "mlx.core.array.split", "mlx.core.array.sqrt", "mlx.core.array.square", "mlx.core.array.sum", "mlx.core.array.tolist", "mlx.core.array.transpose", "mlx.core.array.var", "mlx.core.array_equal", "mlx.core.broadcast_to", "mlx.core.concatenate", "mlx.core.conv1d", "mlx.core.conv2d", "mlx.core.convolve", "mlx.core.cos", "mlx.core.cosh", "mlx.core.default_device", "mlx.core.default_stream", "mlx.core.divide", "mlx.core.equal", "mlx.core.erf", "mlx.core.erfinv", "mlx.core.eval", "mlx.core.exp", "mlx.core.expand_dims", "mlx.core.fft.fft", "mlx.core.fft.fft2", "mlx.core.fft.fftn", "mlx.core.fft.ifft", "mlx.core.fft.ifft2", "mlx.core.fft.ifftn", "mlx.core.fft.irfft", "mlx.core.fft.irfft2", "mlx.core.fft.irfftn", "mlx.core.fft.rfft", "mlx.core.fft.rfft2", "mlx.core.fft.rfftn", "mlx.core.full", "mlx.core.grad", "mlx.core.greater", "mlx.core.greater_equal", "mlx.core.jvp", "mlx.core.less", "mlx.core.less_equal", "mlx.core.load", "mlx.core.log", "mlx.core.log10", "mlx.core.log1p", "mlx.core.log2", "mlx.core.logaddexp", "mlx.core.logical_not", "mlx.core.logsumexp", "mlx.core.matmul", "mlx.core.max", "mlx.core.maximum", "mlx.core.mean", "mlx.core.min", "mlx.core.minimum", "mlx.core.multiply", "mlx.core.negative", "mlx.core.new_stream", "mlx.core.ones", "mlx.core.ones_like", "mlx.core.pad", "mlx.core.partition", "mlx.core.prod", "mlx.core.random.bernoulli", "mlx.core.random.categorical", "mlx.core.random.gumbel", "mlx.core.random.key", "mlx.core.random.normal", "mlx.core.random.randint", "mlx.core.random.seed", "mlx.core.random.split", "mlx.core.random.truncated_normal", "mlx.core.random.uniform", "mlx.core.reciprocal", "mlx.core.reshape", "mlx.core.rsqrt", "mlx.core.save", "mlx.core.savez", "mlx.core.savez_compressed", "mlx.core.set_default_device", "mlx.core.set_default_stream", "mlx.core.sigmoid", "mlx.core.sign", "mlx.core.sin", "mlx.core.sinh", "mlx.core.softmax", "mlx.core.sort", "mlx.core.split", "mlx.core.sqrt", "mlx.core.square", "mlx.core.squeeze", "mlx.core.stop_gradient", "mlx.core.subtract", "mlx.core.sum", "mlx.core.take", "mlx.core.take_along_axis", "mlx.core.tan", "mlx.core.tanh", "mlx.core.transpose", "mlx.core.value_and_grad", "mlx.core.var", "mlx.core.vjp", "mlx.core.vmap", "mlx.core.where", "mlx.core.zeros", "mlx.core.zeros_like", "mlx.nn.Conv1d", "mlx.nn.Conv2d", "mlx.nn.Embedding", "mlx.nn.GELU", "mlx.nn.GroupNorm", "mlx.nn.LayerNorm", "mlx.nn.Linear", "mlx.nn.MultiHeadAttention", "mlx.nn.RMSNorm", "mlx.nn.ReLU", "mlx.nn.RoPE", "mlx.nn.Sequential", "mlx.nn.SiLU", "mlx.nn.value_and_grad", "mlx.optimizers.Adam", "mlx.optimizers.Optimizer", "mlx.optimizers.OptimizerState", "mlx.optimizers.SGD", "mlx.utils.tree_flatten", "mlx.utils.tree_map", "mlx.utils.tree_unflatten", "mlx.nn.gelu", "mlx.nn.gelu_approx", "mlx.nn.gelu_fast_approx", "mlx.nn.relu", "mlx.nn.silu", "Array", "Data Types", "Devices and Streams", "FFT", "Neural Networks", "mlx.nn.Module", "Operations", "Optimizers", "Random", "Transforms", "Tree Utils", "Quick Start Guide", "Using Streams"], "terms": {"mlx": [1, 2, 3, 4, 6, 190, 193, 194, 196, 197], "provid": [1, 3, 89, 153, 167, 179, 190, 191, 198], "open": [1, 15, 122, 126], "flexibl": [1, 5], "which": [1, 3, 4, 5, 6, 15, 33, 73, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 89, 92, 95, 118, 119, 128, 130, 131, 132, 144, 148, 153, 155, 156, 163, 170, 191, 194, 198], "user": [1, 3, 190], "mai": 1, "add": [1, 3, 75, 100, 114, 160, 161], "special": 1, "without": [1, 3, 5, 145, 167, 190, 196], "much": [1, 3], "hassl": 1, "while": [1, 3, 6, 128, 170], "librari": [1, 6, 190], "suppli": 1, "effici": [1, 3, 5, 170, 197], "can": [1, 3, 5, 6, 11, 15, 47, 57, 69, 70, 73, 90, 91, 93, 94, 100, 105, 108, 109, 117, 118, 122, 125, 126, 146, 153, 162, 171, 190, 191, 193, 194, 196, 197, 198], "compos": [1, 5, 190, 197], "ani": [1, 3, 5, 15, 163, 178, 179, 180, 190, 191, 196, 197], "number": [1, 15, 51, 63, 89, 92, 114, 118, 121, 124, 126, 153, 155, 156, 160, 161, 164, 167, 194, 198], "applic": 1, "aris": 1, "case": [1, 3, 78, 81, 82, 84, 85, 86, 87, 103, 128, 144, 197, 198], "where": [1, 4, 153, 156, 160, 161, 163, 164, 165, 168, 172, 181, 182, 183, 185, 191], "new": [1, 4, 60, 128, 152, 167, 179, 191, 193], "function": [1, 2, 3, 4, 5, 13, 26, 71, 72, 73, 89, 92, 103, 135, 153, 155, 156, 163, 171, 173, 179, 181, 182, 183, 190, 191, 193, 194, 196], "highli": [1, 6], "optim": [1, 2, 4, 5, 191], "ar": [1, 2, 3, 4, 5, 6, 13, 15, 59, 60, 64, 77, 78, 80, 81, 83, 84, 86, 87, 89, 92, 103, 114, 115, 117, 118, 119, 122, 125, 126, 131, 132, 144, 148, 153, 155, 156, 160, 161, 164, 165, 167, 178, 179, 190, 191, 196, 197], "need": [1, 3, 4, 5, 59, 190, 191, 194, 197], "For": [1, 3, 6, 180, 191, 194, 197], "you": [1, 3, 5, 6, 194], "design": [1, 2, 5, 194], "your": [1, 3, 6, 191], "own": [1, 6], "link": [1, 6], "top": 1, "core": [1, 2, 3, 4, 190, 191, 193, 197], "we": [1, 2, 3, 4, 162, 171, 190, 194, 196], "inner": 1, "work": [1, 3, 6], "go": [1, 3], "over": [1, 3, 4, 12, 14, 22, 23, 24, 25, 62, 63, 78, 81, 84, 87, 102, 104, 106, 107, 115, 116, 129, 139, 140, 147, 154, 160, 161, 164, 165, 168], "simpl": [1, 3, 4, 162, 190], "learn": [1, 2, 4, 5, 164, 165, 168, 177], "step": [1, 3, 4, 15], "involv": [1, 193], "ad": [1, 2, 6, 191], "let": [1, 2, 3], "s": [1, 2, 3, 4, 35, 44, 77, 78, 80, 81, 83, 84, 86, 87, 89, 95, 106, 118, 153, 154, 156, 173, 175, 190, 191, 193, 194, 197], "sai": [1, 3], "would": [1, 3], "like": [1, 3, 5, 113, 159, 197], "an": [1, 3, 4, 6, 8, 12, 14, 26, 60, 62, 63, 73, 88, 104, 107, 112, 113, 114, 116, 128, 141, 144, 148, 149, 156, 158, 159, 164, 165, 166, 167, 175, 176, 178, 182, 190, 191, 194, 196, 197, 198], "take": [1, 3, 4, 89, 92, 105, 108, 113, 149, 153, 155, 156, 159, 194, 198], "two": [1, 11, 13, 59, 69, 70, 77, 80, 86, 90, 91, 93, 94, 100, 103, 105, 108, 109], "arrai": [1, 3, 4, 5, 8, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 59, 60, 61, 62, 63, 64, 65, 66, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 124, 125, 126, 127, 128, 129, 130, 131, 132, 135, 136, 137, 138, 139, 140, 141, 142, 143, 144, 145, 146, 147, 148, 149, 150, 151, 152, 153, 154, 155, 156, 157, 158, 159, 190, 191, 197], "x": [1, 2, 3, 4, 71, 119, 131, 135, 157, 163, 164, 165, 168, 169, 172, 179, 181, 182, 183, 184, 185, 190, 191, 193, 197], "y": [1, 2, 3, 4, 157, 164, 165, 168, 190, 193], "scale": [1, 3, 167], "them": [1, 3, 190, 191], "both": [1, 11, 69, 70, 90, 91, 93, 94, 100, 105, 108, 109, 118, 146, 193, 197], "some": [1, 2, 3, 4, 191], "coeffic": 1, "alpha": 1, "beta": [1, 164, 165, 174], "respect": [1, 2, 4, 89, 153, 163, 164, 165, 179, 190, 191, 197], "togeth": [1, 4, 179], "get": [1, 2, 4, 63, 120, 176], "z": 1, "well": [1, 3, 167, 190, 191], "veri": [1, 3, 167], "easili": 1, "do": [1, 3, 6, 191], "just": [1, 4], "write": [1, 3, 190], "out": [1, 6], "follow": [1, 3, 4, 5, 6, 15, 64, 174, 177, 182, 183, 194], "import": [1, 2, 3, 4, 6, 131, 153, 178, 179, 180, 190, 191, 197], "mx": [1, 2, 3, 4, 131, 153, 169, 184, 190, 191, 193, 194, 197, 198], "def": [1, 2, 3, 4, 153, 190, 191], "simple_axpbi": 1, "float": [1, 11, 13, 15, 26, 56, 59, 60, 69, 70, 88, 90, 91, 93, 94, 100, 101, 105, 108, 109, 114, 117, 122, 125, 126, 146, 157, 164, 165, 168, 174, 177, 187, 191], "return": [1, 2, 3, 4, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 33, 37, 49, 56, 59, 60, 61, 62, 63, 64, 65, 66, 69, 70, 71, 72, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 124, 125, 126, 127, 128, 129, 135, 136, 137, 138, 139, 140, 141, 142, 143, 144, 145, 146, 147, 148, 149, 150, 151, 152, 153, 154, 155, 156, 157, 158, 159, 173, 178, 179, 180, 190, 191, 196], "thi": [1, 3, 4, 6, 12, 13, 14, 15, 22, 23, 24, 25, 73, 92, 100, 102, 103, 104, 106, 107, 115, 116, 118, 139, 140, 141, 147, 148, 154, 182, 183, 190, 191, 196], "perform": [1, 3, 5, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 103, 139, 148, 164, 190], "leav": [1, 179], "differenti": [1, 5], "howev": [1, 163, 164, 190, 194], "vector": [1, 2, 5, 92, 148, 155, 156, 162, 197], "math": [1, 3], "often": 1, "realiz": 1, "axpbi": 1, "routin": 1, "defin": [1, 2, 3, 4, 6, 176, 178], "same": [1, 3, 6, 59, 60, 63, 64, 82, 85, 86, 87, 89, 92, 114, 118, 155, 157, 164, 190, 191, 194], "realli": 1, "part": 1, "doe": [1, 3, 6, 190], "fast": [1, 163, 183], "so": [1, 3, 6, 89, 153, 193], "decid": [1, 191], "want": [1, 3], "reli": 1, "acceler": 1, "framework": [1, 5], "continu": 1, "impos": 1, "our": [1, 3, 171, 174], "assumpt": 1, "also": [1, 3, 4, 5, 11, 69, 70, 78, 81, 84, 87, 90, 91, 93, 94, 100, 105, 108, 109, 146, 167, 173, 176, 181, 190, 191, 193, 197, 198], "assum": [1, 3, 164, 179, 190], "how": [1, 3, 4, 160, 161, 162, 167, 190], "gradient": [1, 2, 4, 89, 145, 153, 173, 177, 190, 191, 193, 197], "ins": 1, "what": [1, 3], "coincid": 1, "right": [1, 163, 182, 183], "place": [1, 3], "cours": 1, "The": [1, 3, 4, 5, 6, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 33, 35, 44, 49, 56, 59, 60, 61, 62, 63, 64, 65, 66, 69, 70, 71, 72, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 124, 125, 126, 127, 128, 135, 136, 137, 138, 139, 140, 142, 143, 144, 145, 146, 147, 148, 149, 150, 151, 152, 153, 154, 155, 156, 157, 158, 159, 160, 161, 162, 164, 165, 166, 167, 168, 170, 171, 173, 175, 176, 177, 178, 179, 180, 187, 191, 193, 197, 198], "structur": [1, 73], "from": [1, 3, 4, 5, 83, 84, 86, 87, 88, 95, 103, 113, 117, 118, 119, 120, 122, 125, 131, 144, 145, 146, 148, 149, 157, 159, 167, 178, 179, 180, 190, 191, 196, 197], "frontend": 1, "api": 1, "redirect": 1, "when": [1, 3, 5, 160, 161, 190, 191, 194], "appropri": 1, "fallback": 1, "metal": [1, 6], "vjp": [1, 197], "jvp": [1, 197], "In": [1, 3, 4, 103, 164, 174, 179, 190, 191, 196], "one": [1, 3, 6, 56, 63, 75, 98, 103, 118, 144, 146, 191], "sentenc": 1, "comput": [1, 2, 3, 4, 5, 6, 89, 92, 100, 106, 139, 145, 153, 154, 155, 164, 165, 168, 173, 182, 183, 190, 191, 193, 197], "graph": [1, 3, 4, 5, 73, 130], "rule": 1, "evalu": [1, 3, 4, 73, 92, 130, 155, 190, 191, 193, 197], "said": [1, 3], "start": [1, 2, 3, 5, 6, 15, 141], "discuss": 1, "more": [1, 4, 8, 56, 103, 190, 194], "detail": [1, 8, 174, 190], "thei": [1, 2, 3, 64, 171, 191, 196, 197], "c": [1, 3, 160, 161, 187, 197], "scalar": [1, 11, 13, 37, 56, 59, 69, 70, 89, 90, 91, 93, 94, 100, 101, 103, 105, 108, 109, 114, 122, 125, 126, 146, 153, 173, 197], "sum": [1, 2, 11, 102, 139], "elementwis": 1, "numpi": [1, 3, 4, 5, 11, 13, 15, 26, 60, 69, 70, 90, 91, 93, 94, 100, 103, 105, 108, 109, 146, 197], "style": [1, 11, 13, 69, 70, 90, 91, 93, 94, 100, 103, 105, 108, 109, 146], "broadcast": [1, 11, 13, 60, 69, 70, 88, 90, 91, 93, 94, 100, 103, 105, 108, 109, 117, 118, 125, 126, 146, 149, 157, 167], "between": [1, 5], "input": [1, 2, 3, 10, 11, 12, 13, 14, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 59, 60, 61, 62, 63, 64, 65, 66, 69, 70, 71, 72, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 89, 90, 91, 92, 93, 94, 96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 113, 114, 115, 116, 124, 127, 128, 129, 135, 136, 137, 138, 139, 140, 141, 142, 143, 144, 145, 146, 147, 148, 149, 150, 151, 152, 153, 154, 156, 157, 159, 160, 161, 162, 164, 165, 166, 167, 168, 170, 197], "upcast": 1, "const": 1, "factor": 1, "streamordevic": 1, "stream": [1, 5, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 28, 29, 30, 31, 32, 33, 34, 36, 38, 39, 40, 41, 42, 43, 45, 46, 47, 48, 50, 52, 53, 54, 55, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 68, 69, 70, 71, 72, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 90, 91, 93, 94, 95, 96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115, 116, 117, 118, 119, 121, 122, 124, 125, 126, 127, 128, 129, 134, 135, 136, 137, 138, 139, 140, 141, 142, 143, 144, 145, 146, 147, 148, 149, 150, 151, 152, 154, 157, 158, 159], "schedul": 1, "itself": 1, "call": [1, 3, 4, 27, 162, 171, 190, 191, 193], "other": [1, 3, 167, 190, 191, 197], "within": [1, 24], "simplest": 1, "wai": [1, 3, 6, 190], "about": [1, 3, 4], "term": 1, "exist": [1, 3, 191], "auto": 1, "ax": [1, 12, 14, 22, 23, 57, 75, 77, 78, 80, 81, 83, 84, 86, 87, 102, 104, 106, 107, 114, 116, 139, 144, 147, 152, 154], "multipli": 1, "earlier": 1, "goal": 1, "themselv": 1, "contain": [1, 3, 49, 73, 82, 83, 84, 101, 141, 157, 190, 191], "act": 1, "data": [1, 4, 5, 8, 15, 85, 86, 88, 112, 125, 158], "nor": [1, 89, 153], "rather": 1, "easi": [1, 190], "interfac": 1, "block": [1, 3], "A": [1, 3, 5, 6, 49, 59, 89, 92, 102, 103, 117, 118, 119, 121, 122, 125, 126, 141, 153, 155, 156, 164, 165, 168, 171, 173, 174, 178, 179, 180, 183, 190, 191, 193], "It": [1, 3, 6, 89, 153, 175, 190], "creat": [1, 3, 6, 190, 191, 193], "output": [1, 3, 6, 12, 13, 14, 15, 22, 23, 24, 60, 82, 85, 86, 87, 88, 89, 102, 104, 106, 107, 112, 113, 115, 116, 117, 118, 119, 121, 122, 125, 126, 131, 132, 139, 144, 147, 149, 153, 154, 155, 156, 157, 158, 159, 160, 161, 166, 167, 197], "given": [1, 12, 14, 24, 60, 61, 73, 75, 76, 77, 78, 79, 80, 81, 85, 86, 87, 88, 102, 104, 106, 107, 116, 122, 139, 141, 147, 154, 167, 191], "set": [1, 3, 4, 163, 166, 170, 176, 191, 194], "further": [1, 6], "class": [1, 3, 4, 7, 8, 9, 26, 160, 161, 162, 163, 164, 165, 166, 167, 168, 169, 170, 171, 172, 174, 175, 176, 177, 181, 182, 183, 184, 185, 191], "under": 1, "These": [1, 149], "word": 1, "bit": [1, 187, 191], "abstract": 1, "back": [1, 3], "give": [1, 3, 4, 24], "ourselv": 1, "concret": 1, "imag": [1, 161], "public": [1, 190], "explicit": [1, 194], "alpha_": 1, "beta_": 1, "must": [1, 6, 73, 88, 117, 118, 122, 125, 126, 157], "know": [1, 3], "popul": 1, "To": [1, 2, 3, 6, 197], "avoid": 1, "unecessari": [], "alloc": [1, 191], "respons": 1, "space": 1, "void": 1, "eval_cpu": 1, "std": 1, "overrid": 1, "eval_gpu": 1, "jacobian": [1, 92, 155, 197], "product": [1, 92, 103, 116, 155, 167, 197], "primal": [1, 92, 155], "tangent": [1, 20, 21, 92, 150, 151], "int": [1, 3, 4, 7, 9, 11, 12, 14, 15, 22, 23, 24, 25, 26, 29, 30, 31, 32, 40, 41, 42, 43, 45, 49, 52, 55, 56, 58, 59, 60, 61, 62, 63, 69, 70, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 93, 94, 100, 101, 102, 104, 105, 106, 107, 108, 109, 112, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 128, 139, 140, 141, 144, 146, 147, 148, 149, 152, 153, 154, 156, 157, 158, 160, 161, 162, 164, 165, 166, 167, 168, 170, 190, 191], "argnum": [1, 89, 153], "cotan": 1, "accross": [1, 164], "pair": [1, 114, 170], "repres": [1, 3], "axi": [1, 3, 4, 12, 14, 22, 23, 24, 25, 29, 30, 31, 32, 40, 41, 42, 43, 45, 52, 55, 58, 61, 75, 76, 79, 82, 83, 84, 85, 86, 87, 102, 104, 106, 107, 114, 115, 116, 118, 139, 140, 141, 144, 147, 148, 149, 152, 154, 156], "correspond": [1, 12, 14, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 102, 104, 107, 116, 147, 156, 179], "dimens": [1, 3, 12, 14, 22, 23, 44, 49, 56, 63, 75, 83, 84, 86, 87, 102, 103, 104, 106, 107, 116, 118, 124, 147, 149, 152, 154, 160, 161, 164, 165, 167, 168, 170], "vmap": [1, 197], "print": [1, 2, 3, 4, 6, 178, 179, 180, 190, 194, 197], "ostream": 1, "os": [1, 6], "equival": [1, 27, 47, 57, 163], "check": [1, 59], "bool": [1, 11, 12, 14, 22, 23, 26, 29, 30, 31, 32, 40, 41, 42, 43, 45, 55, 56, 58, 59, 60, 69, 70, 73, 88, 90, 91, 93, 94, 100, 101, 102, 104, 105, 106, 107, 108, 109, 114, 116, 117, 122, 125, 126, 130, 146, 147, 154, 157, 160, 161, 164, 165, 166, 170, 191], "is_equival": 1, "privat": 1, "fall": 1, "eval": [1, 2, 3, 4, 190, 191, 193, 197], "deriv": 1, "base": [1, 73, 97, 99, 175, 191, 193, 194], "abov": [1, 3, 6], "demonstr": 1, "treat": [1, 59, 83, 84, 86, 87, 148], "paramet": [1, 2, 3, 4, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 33, 59, 60, 61, 62, 63, 64, 65, 66, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 127, 128, 129, 130, 131, 132, 135, 136, 137, 138, 139, 140, 141, 142, 143, 144, 145, 146, 147, 148, 149, 150, 151, 152, 153, 154, 155, 156, 157, 158, 159, 160, 161, 162, 163, 164, 165, 166, 167, 168, 170, 171, 173, 175, 177, 178, 179, 180, 191, 193], "produc": [1, 167], "through": [1, 145], "construct": [1, 4, 88, 112, 158], "its": [1, 6, 103, 115, 124, 173, 180, 190], "type": [1, 5, 7, 8, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 33, 49, 56, 59, 60, 61, 62, 63, 64, 65, 66, 69, 70, 71, 72, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 124, 125, 126, 127, 128, 129, 135, 136, 137, 138, 139, 140, 141, 142, 143, 144, 145, 146, 147, 148, 149, 150, 151, 152, 153, 154, 155, 156, 157, 158, 159, 175, 178, 190], "shape": [1, 3, 4, 47, 59, 60, 62, 63, 76, 79, 82, 85, 86, 87, 88, 92, 103, 112, 113, 117, 118, 119, 121, 122, 125, 126, 128, 149, 155, 157, 158, 159, 160, 161, 190, 193, 197], "pass": [1, 3, 4, 47, 57, 114, 153, 171, 173, 178, 190, 191], "re": [1, 4], "now": [1, 3], "promot": 1, "dtype": [1, 3, 15, 26, 33, 56, 88, 112, 119, 121, 122, 125, 126, 158, 187, 197], "promoted_dtyp": 1, "promote_typ": 1, "float32": [1, 15, 112, 119, 121, 125, 126, 158, 187, 197], "non": [1, 6, 191], "point": [1, 2, 3, 6, 187], "out_dtyp": 1, "is_floating_point": 1, "cast": [1, 33, 85, 86, 87, 191], "up": [1, 3], "determin": 1, "x_cast": 1, "astyp": [1, 3, 191], "y_cast": 1, "broadcasted_input": 1, "broadcast_arrai": 1, "out_shap": 1, "0": [1, 2, 3, 4, 6, 7, 15, 52, 58, 61, 62, 63, 89, 114, 117, 126, 141, 153, 154, 156, 160, 161, 163, 164, 165, 169, 174, 177, 178, 182, 183, 184, 190, 191, 194, 197], "unique_ptr": 1, "make_uniqu": 1, "to_stream": 1, "handl": [1, 190], "resolv": 1, "No": [1, 3], "happen": [1, 3, 193], "alon": 1, "effect": 1, "onli": [1, 3, 5, 6, 59, 62, 63, 187, 190, 191], "execut": [1, 6], "depend": [1, 2, 56], "devic": [1, 5, 6, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 28, 29, 30, 31, 32, 33, 34, 36, 38, 39, 40, 41, 42, 43, 45, 46, 47, 48, 50, 52, 53, 54, 55, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 90, 91, 93, 94, 95, 96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115, 116, 117, 118, 119, 121, 122, 124, 125, 126, 127, 128, 129, 133, 135, 136, 137, 138, 139, 140, 141, 142, 143, 144, 145, 146, 147, 148, 149, 150, 151, 152, 154, 157, 158, 159, 198], "specifi": [1, 15, 33, 63, 83, 84, 88, 89, 112, 118, 148, 149, 152, 153, 156, 158], "memori": [1, 5, 191], "ha": [1, 3, 4, 5, 56, 82, 83, 85, 86, 87, 89, 118, 191, 193, 197], "been": [1, 3], "try": 1, "naiv": 1, "gener": [1, 2, 15, 83, 84, 117, 121, 122, 125, 126, 194, 198], "version": [1, 6, 100, 102, 139, 156, 194], "declar": 1, "member": [1, 190, 191], "method": [1, 3, 7, 8, 9, 26, 174, 175, 176, 177], "each": [1, 49, 73, 103, 114, 118, 131, 132, 141, 152, 156, 157, 162, 164, 194], "element": [1, 10, 11, 16, 17, 18, 19, 20, 21, 24, 65, 66, 69, 70, 71, 72, 74, 90, 91, 93, 94, 96, 97, 98, 99, 100, 101, 105, 108, 109, 110, 115, 127, 129, 135, 136, 137, 138, 142, 143, 146, 148, 150, 151, 153, 157, 170, 172, 185], "find": [1, 2, 6], "pointwis": 1, "captur": [1, 190], "templat": 1, "axpby_impl": 1, "typenam": 1, "t": [1, 3, 71, 153, 174, 177, 190], "readi": 1, "fill": [1, 88, 113, 159], "malloc_or_wait": 1, "synchron": 1, "avail": [1, 2, 3, 4, 6, 8, 187], "There": [1, 190], "wait": [1, 3], "here": [1, 3], "request": 1, "pressur": 1, "condit": [1, 157], "set_data": 1, "nbyte": 1, "collect": [1, 176, 179, 196], "pointer": 1, "x_ptr": 1, "y_ptr": 1, "out_ptr": 1, "relev": 1, "static_cast": 1, "size_t": 1, "out_idx": 1, "size": [1, 3, 4, 49, 63, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 118, 128, 141, 144, 160, 161, 162], "map": [1, 4, 95, 162, 179, 191], "linear": [1, 3, 4, 5, 163, 169, 172, 179, 181, 182, 183, 184, 185, 190, 191], "indic": [1, 13, 22, 23, 24, 25, 73, 89, 141, 148, 149, 153], "offset": [1, 3], "x_offset": 1, "elem_to_loc": 1, "stride": [1, 62, 63, 160, 161, 170], "y_offset": 1, "contigu": 1, "regularli": 1, "default": [1, 6, 12, 14, 15, 22, 23, 24, 25, 59, 61, 62, 63, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 89, 102, 104, 106, 107, 112, 115, 116, 117, 118, 119, 121, 122, 124, 125, 126, 128, 130, 140, 141, 144, 147, 152, 153, 154, 156, 158, 160, 161, 167, 170, 176, 187, 191, 194, 196, 198], "row": 1, "major": 1, "henc": 1, "doesn": [1, 190], "additon": 1, "abl": 1, "all": [1, 4, 6, 13, 24, 63, 73, 78, 81, 84, 87, 103, 114, 115, 144, 167, 175, 190, 191, 194, 197, 198], "incom": 1, "accordingli": 1, "dispatch": 1, "float16": [1, 187, 191], "bfloat16": 1, "complex64": 1, "throw": 1, "error": [1, 71, 72, 141, 163, 181, 182, 183], "encount": 1, "unexpect": [1, 15], "regist": [1, 4], "op": 1, "contruct": 1, "assert": 1, "2": [1, 2, 3, 4, 63, 71, 77, 80, 82, 83, 84, 85, 86, 87, 99, 103, 124, 161, 163, 168, 174, 182, 187, 190, 191, 197], "1": [1, 3, 4, 15, 24, 25, 62, 63, 76, 77, 79, 80, 82, 83, 84, 85, 86, 87, 103, 115, 118, 126, 135, 140, 148, 153, 160, 161, 163, 164, 165, 168, 170, 174, 177, 182, 183, 187, 191, 193, 197], "correct": [1, 174], "els": [1, 3, 190, 191], "float16_t": 1, "bfloat16_t": 1, "complex64_t": 1, "runtime_error": 1, "support": [1, 3, 5, 6, 13, 62, 63, 103], "have": [1, 3, 6, 59, 83, 84, 86, 87, 103, 118, 167, 171, 178, 196], "rememb": 1, "3": [1, 3, 6, 194, 197], "complic": 1, "keep": [1, 12, 14, 22, 23, 102, 104, 106, 107, 116, 147, 154, 190, 191], "mind": [1, 3], "half": [1, 15, 122, 126, 170], "precis": [1, 3, 163, 190], "direct": [1, 3, 191], "fix": [1, 3], "possibl": [1, 3, 103, 141, 162], "due": 1, "transpos": [1, 3, 27], "aren": 1, "guarante": 1, "fit": 1, "requir": [1, 3, 190], "column": 1, "inplac": 1, "expect": [1, 3, 160, 161, 167], "answer": 1, "copi": [1, 3, 5, 115, 140], "simpli": [1, 3, 6, 169, 184, 191], "catlas_saxpbi": 1, "axpby_impl_acceler": 1, "first": [1, 2, 3, 4, 6, 89, 103, 115, 124, 153, 164, 174, 178, 190], "mode": [1, 64], "i": [1, 3, 92, 160, 161, 190], "e": [1, 4, 6, 71, 92, 135, 160, 161, 164, 165, 168, 190, 193, 198], "match": [1, 6, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87], "transposit": 1, "data_s": 1, "items": 1, "flag": 1, "copy_inplac": 1, "copytyp": 1, "n": [1, 3, 26, 62, 63, 76, 78, 79, 81, 82, 85, 87, 154, 160, 161], "incx": 1, "inci": 1, "great": 1, "But": 1, "criteria": 1, "luckili": 1, "alwai": [1, 178], "With": 1, "final": [1, 2, 3, 4], "singl": [1, 4, 73, 92, 114, 155], "row_contigu": 1, "col_contigu": 1, "common": 1, "hit": 1, "mileston": 1, "enough": 1, "run": [1, 3, 4, 5, 6, 191, 198], "If": [1, 3, 6, 12, 14, 15, 22, 23, 24, 25, 56, 59, 61, 64, 73, 85, 86, 87, 88, 89, 102, 103, 104, 106, 107, 112, 114, 115, 116, 118, 139, 140, 141, 147, 148, 149, 153, 154, 156, 158, 160, 161, 164, 165, 166, 167, 170, 171, 179, 191, 198], "plan": 1, "stop": [1, 3, 15, 145], "enjoi": 1, "speed": 1, "appl": [1, 3, 5, 6], "silicon": [1, 3, 5, 6], "address": 1, "shade": 1, "languag": [1, 187], "kernel": [1, 62, 63], "written": 1, "help": [1, 3], "resourc": 1, "walkthrough": 1, "pipelin": 1, "specif": [1, 6], "cpp": 1, "algorithm": 1, "launch": 1, "exactli": [1, 3], "mani": [1, 141, 160, 161, 162, 167], "thread": 1, "pick": 1, "updat": [1, 2, 3, 4, 177, 179, 191, 193], "assign": [1, 191], "axpby_gener": 1, "buffer": [1, 26], "constant": [1, 3, 6, 114, 164, 165, 168], "4": [1, 3, 131, 187, 197], "5": [1, 2, 3, 6, 117], "x_stride": 1, "6": [1, 3, 131, 182, 183, 197], "y_stride": 1, "7": [1, 3], "ndim": 1, "8": [1, 3, 6, 187, 197], "uint": 1, "index": [1, 7, 9, 24, 75, 89, 115, 148, 149, 153], "thread_position_in_grid": 1, "convert": [1, 56, 197], "instanti": [1, 4], "uniqu": [1, 194], "host": 1, "name": [1, 95, 131, 132, 164, 176, 190, 191], "identifi": [1, 178, 196], "instantiate_axpbi": 1, "type_nam": 1, "host_nam": 1, "axpby_general_": 1, "bflot16": 1, "compil": [1, 6], "mlx_ext": 1, "metallib": [1, 6], "see": [1, 3, 4, 8, 28, 29, 30, 31, 32, 34, 36, 38, 39, 40, 41, 42, 43, 45, 46, 47, 48, 50, 52, 53, 54, 55, 57, 58, 163, 181, 182, 183], "later": [1, 6], "co": 1, "locat": [1, 191], "share": [1, 5], "register_librari": 1, "potenti": 1, "path": [1, 6, 131, 132], "tri": 1, "load": [1, 4, 191], "hasn": 1, "alreadi": [1, 3], "static": [1, 6], "object": [1, 8, 26, 37, 56, 95, 130, 131, 132, 156, 178, 196], "why": [1, 3], "packag": [1, 2, 4], "process": [1, 3, 64, 162, 179, 196], "logic": [1, 101], "grid": 1, "shown": 1, "below": [1, 187], "prepar": [1, 3], "carri": 1, "should": [1, 2, 3, 4, 6, 73, 92, 149, 153, 155, 160, 161, 167, 171, 178, 190, 191, 196, 198], "d": [1, 3, 103, 148, 174, 180], "ostringstream": 1, "kname": 1, "axpby_": 1, "general_": 1, "type_to_nam": 1, "make": [1, 3, 4, 6, 103, 190, 197], "sure": [1, 3, 6, 190], "look": [1, 3], "folder": 1, "get_colocated_mtllib_path": 1, "get_kernel": 1, "str": [1, 64, 89, 95, 130, 131, 132, 153, 178, 180, 191], "encod": [1, 170], "compute_encod": 1, "get_command_encod": 1, "setcomputepipelinest": 1, "those": [1, 3, 190], "decelar": 1, "nelem": 1, "set_array_buff": 1, "setbyt": 1, "sizeof": 1, "threadgroup": 1, "higher": 1, "than": [1, 3, 56, 64, 90, 91, 93, 94, 103, 170, 179], "max": [1, 105], "allow": [1, 175, 190, 191, 197], "tgp_size": 1, "min": [1, 108], "maxtotalthreadsperthreadgroup": 1, "3d": 1, "mtl": 1, "group_dim": 1, "grid_dim": 1, "divd": 1, "among": 1, "dispatchthread": 1, "few": [1, 3, 4, 5, 197], "thing": [1, 3], "note": [1, 3, 6, 13, 62, 63, 83, 84, 118, 190], "befor": [1, 3, 6, 24, 115, 130, 191], "move": 1, "track": [1, 190], "activ": [1, 190], "command": [1, 6], "instead": [1, 190], "end_encod": 1, "end": 1, "until": [1, 197], "limit": 1, "flush": 1, "enqueu": 1, "commit": 1, "associ": [1, 131, 132], "suggest": 1, "deeper": 1, "dive": 1, "studi": 1, "come": [1, 3], "far": [1, 193], "built": [1, 6], "includ": [1, 191, 197, 198], "forward": [1, 153], "diff": 1, "push": 1, "along": [1, 22, 23, 61, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 139, 141, 148, 149], "primtiv": 1, "similarli": [1, 6, 103], "scale_arr": 1, "contribut": 1, "tangent_x": 1, "tangent_i": 1, "revers": [1, 152], "arg": [1, 3, 8, 26, 47, 57, 73, 131, 132], "push_back": 1, "fulli": [1, 5], "primitv": 1, "overal": 1, "directori": [1, 3, 6], "extens": [1, 187], "h": [1, 62, 63, 161], "mlx_sample_extens": 1, "__init__": [1, 3, 4, 7, 8, 9, 26, 190, 191], "py": [1, 3], "cmakelist": 1, "txt": 1, "setup": [1, 2, 4], "strucutr": 1, "hold": [1, 3, 8, 175], "instal": 1, "pybind11": [1, 6], "sinc": [1, 3, 4, 191], "compon": [1, 3], "etc": [1, 190], "becom": 1, "pybind11_modul": 1, "m": [1, 6], "doc": [1, 4], "sampl": [1, 2, 3, 117, 118, 119, 122, 125, 126, 194], "_a": 1, "pos_onli": 1, "kw_onli": 1, "none": [1, 3, 7, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 28, 29, 30, 31, 32, 33, 34, 36, 38, 39, 40, 41, 42, 43, 45, 46, 47, 48, 50, 52, 53, 54, 55, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 93, 94, 95, 96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 112, 113, 114, 115, 116, 117, 118, 119, 121, 122, 123, 124, 125, 126, 127, 128, 129, 130, 131, 132, 133, 134, 135, 136, 137, 138, 139, 140, 141, 142, 143, 144, 145, 146, 147, 148, 149, 150, 151, 152, 153, 154, 156, 157, 158, 159, 163, 167, 178, 191], "r": [1, 3, 153], "pbdoc": 1, "most": [1, 118, 190], "complex": [1, 11, 26, 59, 60, 69, 70, 83, 84, 85, 86, 87, 88, 90, 91, 93, 94, 100, 101, 105, 108, 109, 114, 117, 122, 125, 126, 146, 157, 178, 190], "addit": [1, 3, 11, 164, 165, 167, 168, 191], "bell": 1, "whistl": 1, "liter": 1, "string": 1, "modul": [1, 3, 4, 171, 173, 196], "ensur": 1, "caster": 1, "find_packag": 1, "config": 1, "add_librari": 1, "sourc": [1, 152], "target_sourc": 1, "cmake_current_list_dir": 1, "header": 1, "target_include_directori": 1, "target_link_librari": 1, "attach": 1, "conveni": [1, 4], "mlx_build_metallib": 1, "target": [1, 153], "destin": 1, "automat": [1, 5, 197], "practic": 1, "mlx_build_met": [1, 6], "mlx_ext_metallib": 1, "titl": 1, "include_dir": 1, "project_source_dir": 1, "mlx_include_dir": 1, "output_directori": 1, "cmake_library_output_directori": 1, "add_depend": 1, "endif": 1, "pybind11_add_modul": 1, "build_shared_lib": 1, "target_link_opt": 1, "wl": 1, "rpath": 1, "loader_path": 1, "onc": 1, "describ": 1, "util": [1, 3, 5, 131], "__name__": [1, 3], "__main__": [1, 3], "descript": [1, 3, 187], "ext_modul": 1, "cmakeextens": 1, "cmdclass": 1, "build_ext": 1, "cmakebuild": 1, "package_dir": 1, "package_data": 1, "dylib": 1, "zip_saf": 1, "fals": [1, 3, 12, 14, 22, 23, 29, 30, 31, 32, 40, 41, 42, 43, 45, 55, 58, 59, 73, 102, 104, 106, 107, 116, 147, 154, 157, 164, 166, 170, 178, 187, 191], "python_requir": 1, "even": [1, 3], "though": [1, 3], "j8": 1, "libmlx_ext": 1, "cpython": 1, "3x": 1, "darwin": 1, "pip": [1, 6], "after": [1, 3, 4, 24, 115, 164, 165, 167], "plai": [1, 3], "ones": [1, 3, 113, 131, 191], "b": [1, 3, 11, 13, 59, 69, 70, 90, 91, 93, 94, 100, 103, 105, 108, 109, 146, 153, 197], "f": [1, 2, 4, 190], "item": [1, 2, 3, 4, 179, 197], "true": [1, 2, 3, 59, 130, 139, 157, 160, 161, 164, 165, 166, 170, 178, 187, 191], "quick": [1, 5], "benchmark": 1, "compar": [1, 59], "time": [1, 3, 6, 190], "set_default_devic": 1, "256": [1, 4], "512": [1, 3], "random": [1, 2, 3, 4, 5, 198], "normal": [1, 2, 3, 125, 164, 165, 168, 176], "bench": 1, "warm": 1, "rang": [1, 2, 3, 4, 6, 15, 182, 183, 193, 194], "100": [1, 2, 3], "5000": 1, "simple_tim": 1, "custom_tim": 1, "3f": [1, 4], "custom": 1, "114": 1, "109": 1, "modest": 1, "improv": [1, 3], "awai": [1, 3], "good": [1, 6], "nn": [1, 3, 4, 131, 179, 190, 193], "grad": [1, 2, 4, 153, 193, 197], "simplifi": 1, "full": [1, 4, 47, 57, 64, 139, 191], "implement": [2, 4, 162, 167, 170, 171, 174, 175, 176, 191], "basic": 2, "model": [2, 4, 5, 131, 167, 173, 179, 190, 191, 193], "problem": [2, 4, 190], "metadata": 2, "num_featur": 2, "num_exampl": 2, "1_000": 2, "num_it": 2, "10_000": 2, "iter": [2, 4, 179, 194], "sgd": [2, 4, 193], "lr": 2, "01": 2, "rate": 2, "ll": [2, 4], "synthet": 2, "dataset": 2, "matrix": [2, 103], "ground": [2, 3], "truth": 2, "w_star": 2, "valu": [2, 3, 10, 15, 22, 23, 37, 56, 59, 73, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 114, 117, 118, 119, 121, 122, 125, 126, 148, 149, 153, 156, 167, 173, 176, 178, 179, 187, 191], "gaussian": [2, 163, 181, 182, 183], "nois": 2, "exampl": [2, 3, 4, 15, 148, 193, 194, 197], "noisi": 2, "label": 2, "ep": [2, 164, 165, 168, 174], "1e": [2, 4, 13, 164, 165, 168, 174], "us": [2, 3, 4, 5, 6, 15, 103, 128, 162, 163, 166, 167, 175, 178, 182, 183, 190, 191, 193, 194, 196, 197], "weight": [2, 62, 63, 179, 190, 191], "squar": [2, 3, 129, 142, 153, 168, 179, 190], "loss": [2, 4, 153, 190, 193], "loss_fn": [2, 4, 193], "w": [2, 63, 153, 161, 177], "mean": [2, 3, 4, 153, 164, 168, 190, 191], "grad_fn": 2, "initi": [2, 3, 164, 165, 168, 190, 191], "randomli": [2, 3], "Then": [2, 6], "repeatedli": 2, "_": [2, 3, 190, 194], "verifi": 2, "close": [2, 5, 13], "error_norm": 2, "5f": 2, "someth": [2, 3], "00005": 2, "00364": 2, "complet": [2, 3, 191], "logist": [2, 135, 172, 182, 183, 185], "github": [2, 4, 6], "repo": [2, 4, 6], "enabl": [3, 73], "larg": [3, 190], "ish": 3, "transform": [3, 5, 73, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 164, 165, 166, 173, 190, 191], "compromis": 3, "eas": 3, "llama": 3, "famili": 3, "less": [3, 24, 94, 115, 170], "200": 3, "line": 3, "python": [3, 37, 49, 56, 73, 178, 179, 180, 191, 196], "neural": [3, 5, 162, 191], "network": [3, 5, 162, 191], "build": [3, 5, 191], "concis": 3, "architectur": 3, "notabl": 3, "rope": 3, "posit": [3, 24, 89, 115, 153, 160, 161, 167, 170, 179, 190], "option": [3, 12, 14, 15, 22, 23, 24, 25, 26, 31, 32, 61, 62, 63, 64, 73, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 102, 104, 106, 107, 112, 114, 115, 116, 117, 118, 119, 121, 122, 124, 125, 126, 128, 130, 139, 140, 141, 144, 147, 148, 149, 152, 153, 154, 156, 158, 160, 161, 167, 178, 191, 194, 198], "kei": [3, 117, 118, 119, 121, 122, 124, 125, 126, 167, 176, 178, 179, 191, 194, 196], "cach": 3, "concaten": 3, "project": [3, 167], "llamaattent": 3, "self": [3, 4, 7, 9, 26, 27, 28, 29, 30, 31, 32, 33, 34, 36, 37, 38, 39, 40, 41, 42, 43, 45, 46, 47, 48, 50, 52, 53, 54, 55, 56, 57, 58, 190, 191], "dim": [3, 162, 164, 165, 167, 168, 170], "num_head": [3, 167], "super": [3, 4, 190, 191], "tradit": [3, 170], "query_proj": 3, "bia": [3, 160, 161, 166, 174, 179, 191], "key_proj": 3, "value_proj": 3, "out_proj": [3, 191], "__call__": [3, 4, 190, 191], "queri": [3, 167], "mask": [3, 167], "extract": [3, 190, 191], "l": [3, 4, 160, 190], "reshap": 3, "combin": 3, "key_cach": 3, "value_cach": 3, "sqrt": [3, 71, 164, 165, 168, 174], "score": 3, "softmax": 3, "values_hat": 3, "rm": 3, "swiglu": 3, "rmsnorm": 3, "llamaencoderlay": 3, "mlp_dim": 3, "norm1": 3, "norm2": 3, "linear1": 3, "linear2": 3, "linear3": 3, "sigmoid": [3, 172, 182, 183, 185], "instanc": [3, 171, 180, 190, 191], "embed": 3, "emb": [3, 162], "token": [3, 162], "num_lay": [3, 4, 193], "vocab_s": 3, "norm": [3, 164], "multiheadattent": 3, "create_additive_causal_mask": 3, "list": [3, 8, 12, 14, 26, 29, 30, 40, 41, 42, 43, 45, 49, 52, 55, 56, 58, 60, 61, 73, 75, 77, 78, 80, 81, 83, 84, 86, 87, 88, 89, 92, 102, 104, 106, 107, 112, 114, 116, 117, 118, 119, 121, 122, 125, 126, 128, 139, 141, 144, 147, 152, 153, 154, 155, 158, 174, 178, 180, 190, 191, 196], "still": [3, 6], "consid": [3, 13, 59, 164, 178, 196], "train": [3, 4, 191], "ignor": 3, "whatsoev": 3, "rest": [3, 170, 179], "subsect": 3, "prompt": 3, "autoregress": 3, "yield": [3, 4, 194], "temp": 3, "causal": 3, "save": [3, 95, 131, 132, 191], "append": [3, 103], "store": 3, "per": [3, 4, 164, 165, 168, 175], "care": 3, "last": [3, 25, 56, 78, 81, 83, 84, 86, 87, 103, 118, 140, 160, 161, 164], "logit": [3, 118], "next": [3, 4], "categor": 3, "lazili": [3, 190], "noth": [3, 190], "yet": [3, 190, 191, 197], "forc": [3, 4, 190, 197], "choos": [3, 170], "pars": 3, "feed": 3, "loop": [3, 4], "unsqueez": 3, "sequenc": [3, 160, 194], "length": [3, 144, 160], "len": [3, 78, 81, 84, 87], "overwrit": 3, "discard": [3, 178], "old": 3, "moment": [3, 174], "anymor": 3, "everyth": 3, "small": [3, 164, 165, 168], "10": [3, 4, 97, 131, 179, 190], "12": 3, "8192": 3, "1024": 3, "actual": [3, 15, 191], "materi": [3, 5], "could": [3, 190], "20_000": 3, "machin": [3, 5, 6], "8gb": 3, "ram": 3, "32": [3, 4, 187], "44": 3, "doubl": 3, "bracket": 3, "becaus": [3, 190], "batch": [3, 103, 160, 161, 167], "zip": [3, 4], "haven": 3, "anyth": [3, 153], "result": [3, 15, 56, 95, 103, 157, 179], "similar": [3, 167, 191], "runtim": 3, "section": [3, 141], "access": [3, 37, 190, 191], "origin": [3, 174], "sentencepiec": 3, "pytorch": [3, 5, 164], "compat": [3, 118], "npz": [3, 95, 131, 132, 191], "file": [3, 6, 95, 130, 131, 132, 191], "directli": 3, "argpars": 3, "itertool": [3, 179], "starmap": [3, 179], "np": [3, 4, 197], "torch": 3, "map_torch_to_mlx": 3, "tok_embed": 3, "elif": 3, "replac": [3, 191], "attention_norm": 3, "ffn_norm": 3, "wq": 3, "wk": 3, "wv": 3, "wo": 3, "w1": 3, "w2": 3, "w3": 3, "ffn": 3, "separ": [3, 47, 57, 164], "submodul": [3, 4, 190, 191], "feed_forward": 3, "parser": 3, "argumentpars": 3, "add_argu": 3, "torch_weight": 3, "output_fil": 3, "parse_arg": 3, "state": [3, 4, 175, 176, 190, 193, 194], "savez": 3, "k": [3, 191], "v": [3, 64, 191], "left": [3, 163, 170, 182, 183], "disk": 3, "text": 3, "format": [3, 95, 130, 131, 132], "oper": [3, 5, 33, 139, 145, 149, 190, 197, 198], "dictionari": [3, 175, 176, 178, 190, 191, 196], "represent": [3, 178, 180], "tree_unflatten": 3, "helper": 3, "weight_fil": 3, "incur": 3, "sever": [3, 62, 63, 131, 132], "unnecessari": [1, 3], "futur": 3, "pth": 3, "current": [3, 5, 6, 62, 63, 190], "around": 3, "m1": 3, "ultra": 3, "7b": 3, "me": 3, "ishmael": 3, "year": 3, "ago": 3, "never": 3, "long": 3, "info": 3, "247": 3, "press": 3, "enter": 3, "littl": 3, "monei": 3, "my": [3, 6], "purs": 3, "greater": [3, 24, 91, 115], "consequ": 3, "walk": 3, "down": 3, "gower": 3, "street": 3, "afternoon": 3, "heavi": 3, "rain": 3, "saw": 3, "off": [3, 6], "man": 3, "rag": 3, "who": 3, "sat": 3, "upon": [3, 179], "hi": 3, "bundl": 3, "hard": 3, "wet": 3, "he": 3, "were": 3, "cry": 3, "watch": 3, "him": 3, "observ": 3, "numer": [3, 100, 102, 139, 164, 165, 168], "crowd": 3, "wa": [3, 176], "hurri": 3, "437": 3, "330": 3, "second": [3, 103, 153, 174], "spent": 3, "amount": 3, "39": 3, "ms": 3, "By": 3, "bigger": 3, "remain": [3, 153], "almost": 3, "nobodi": 3, "took": 3, "least": 3, "notic": 3, "distanc": 3, "had": 3, "doubt": 3, "minut": 3, "straight": 3, "slowli": 3, "rais": [3, 141], "ey": 3, "speak": 3, "resum": 3, "postur": 3, "stood": 3, "feel": 3, "pain": 3, "heart": 3, "smile": 3, "face": 3, "am": 3, "someon": 3, "three": 3, "quarter": 3, "hour": 3, "made": 3, "immedi": [3, 191], "repli": 3, "again": [3, 190], "hand": 3, "did": 3, "accustom": 3, "thu": [3, 190], "question": 3, "reason": 3, "tell": 3, "understand": 3, "579": 3, "690": 3, "num": [3, 124], "500": 3, "628": 3, "went": 3, "nervou": 3, "trembl": 3, "told": 3, "And": 3, "perhap": 3, "surpris": 3, "matter": [3, 190], "shall": 3, "anyhow": 3, "friend": 3, "ye": 3, "slight": 3, "kind": 3, "longer": [3, 64], "soon": 3, "unless": [3, 191], "unlik": [3, 13], "strang": 3, "amus": 3, "That": 3, "secret": 3, "disappoint": 3, "mine": 3, "cannot": 3, "happi": 3, "ask": 3, "Is": 3, "shop": 3, "bui": 3, "food": 3, "633": 3, "21": 3, "475": 3, "su": 3, "j": [3, 6, 174], "lu": 3, "pan": 3, "murtadha": 3, "wen": 3, "liu": 3, "2021": 3, "roform": 3, "enhanc": 3, "rotari": [3, 170], "arxiv": [3, 164, 165, 168, 170], "preprint": 3, "2104": [3, 170], "09864": [3, 170], "zhang": 3, "sennrich": 3, "2019": 3, "root": [3, 129, 142, 168], "advanc": 3, "inform": [3, 4, 163, 167], "system": 3, "shazeer": 3, "2020": 3, "glu": 3, "variant": 3, "2002": 3, "05202": 3, "classifi": 4, "mnist": 4, "As": [4, 148], "mlp": [4, 190, 193], "inherit": [4, 196], "standard": [4, 37, 56, 103, 119, 197], "idiom": 4, "input_dim": [4, 166], "hidden_dim": [4, 191, 193], "output_dim": [4, 166], "layer_s": 4, "idim": 4, "odim": 4, "maximum": [4, 22, 169, 182, 183, 184, 190, 191], "cross": 4, "entropi": 4, "sub": [4, 124], "commonli": [4, 191], "cross_entropi": 4, "accuraci": 4, "valid": [4, 64, 156, 178, 191, 196], "eval_fn": 4, "argmax": 4, "num_class": [4, 193], "batch_siz": [4, 193], "num_epoch": [4, 193], "learning_r": [4, 174, 177, 193], "train_imag": [4, 193], "train_label": [4, 193], "test_imag": 4, "test_label": 4, "shuffl": 4, "minibatch": 4, "batch_iter": [4, 193], "perm": 4, "permut": 4, "id": 4, "put": 4, "trainabl": [4, 173, 190, 191], "loss_and_grad_fn": [4, 193], "value_and_grad": [4, 190, 191, 193, 197], "epoch": 4, "test": [4, 6], "confus": 4, "decent": 4, "95": 4, "except": [5, 82, 83, 85, 86, 87, 164], "featur": [5, 62, 63, 164, 165, 166, 168, 170], "main": [5, 179, 190], "differ": [5, 146], "lazi": [5, 191, 197], "multi": [5, 160, 161], "cpu": 5, "gpu": 5, "strongli": [], "inspir": 5, "jax": [5, 194], "arrayfir": 5, "noteabl": 5, "unifi": 5, "live": 5, "guid": 5, "regress": 5, "layer": [5, 164, 165, 166, 171, 191], "perceptron": 5, "llm": 5, "infer": [5, 88], "fft": 5, "tree": [5, 73, 89, 153, 156, 175, 178, 179, 180], "develop": [5, 6], "document": [5, 47, 57], "17": 6, "g": [6, 177, 190, 198], "clang": 6, "cmake": 6, "24": 6, "clone": 6, "git": 6, "com": 6, "ml": 6, "explor": 6, "cd": 6, "brew": 6, "conda": 6, "global": [6, 123, 194], "env": 6, "cmake_build_parallel_level": 6, "edit": 6, "unittest": 6, "discov": 6, "mkdir": 6, "p": [6, 117, 174], "either": [6, 11, 47, 56, 57, 69, 70, 90, 91, 93, 94, 100, 103, 105, 108, 109, 146, 153, 171], "libmlx": 6, "preprocessor": 6, "metal_path": 6, "mlx_build_test": 6, "ON": 6, "mlx_build_exampl": 6, "mlx_build_benchmark": 6, "mlx_build_python_bind": 6, "devicetyp": 7, "attribut": [7, 8, 9, 26], "kwarg": [8, 26, 131, 132, 198], "union": [10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 28, 29, 30, 31, 32, 33, 34, 36, 38, 39, 40, 41, 42, 43, 45, 46, 47, 48, 50, 52, 53, 54, 55, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 69, 70, 71, 72, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 93, 94, 95, 96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 112, 113, 114, 115, 116, 117, 118, 119, 121, 122, 124, 125, 126, 127, 128, 129, 135, 136, 137, 138, 139, 140, 141, 142, 143, 144, 145, 146, 147, 148, 149, 150, 151, 152, 153, 154, 157, 158, 159, 161, 191], "wise": [10, 11, 16, 17, 18, 19, 20, 21, 65, 66, 69, 70, 71, 72, 74, 90, 91, 93, 94, 96, 97, 98, 99, 100, 101, 105, 108, 109, 110, 127, 129, 135, 136, 137, 138, 142, 143, 146, 150, 151, 172, 185], "absolut": [10, 13, 182, 183], "semant": [11, 60, 69, 70, 90, 91, 93, 94, 100, 103, 105, 108, 109, 146], "keepdim": [12, 14, 22, 23, 29, 30, 31, 32, 40, 41, 42, 43, 45, 55, 58, 102, 104, 106, 107, 116, 139, 147, 154], "reduct": [12, 14, 102, 104, 107, 116], "reduc": [12, 14, 22, 23, 102, 104, 106, 107, 116, 147, 154], "unspecifi": [12, 14, 15, 22, 23, 24, 25, 61, 88, 102, 104, 106, 107, 112, 115, 116, 139, 140, 147, 148, 154, 158, 198], "entir": [12, 14, 22, 23, 102, 104, 106, 107, 116, 147, 154], "singleton": [12, 14, 22, 23, 102, 103, 104, 106, 107, 116, 147, 154], "rtol": 13, "05": [13, 164, 165, 168], "atol": 13, "08": [13, 174], "approxim": [13, 163, 181, 182, 183], "comparison": [13, 70, 90, 91, 93, 94], "equal": [13, 24, 59, 91, 94, 115, 122, 141], "ab": [13, 153, 164, 165, 168, 170], "array_equ": 13, "rel": 13, "toler": 13, "boolean": [13, 59, 101, 187], "interv": [15, 122, 126], "increment": 15, "otherwis": [15, 178, 191], "int32": [15, 122, 187, 197], "convent": [15, 64], "lead": 15, "fraction": 15, "integr": [15, 148], "invers": [16, 17, 18, 19, 20, 21, 72, 79, 80, 81, 82, 83, 84], "cosin": [16, 17, 65, 66], "hyperbol": [17, 19, 21, 66, 138, 151], "sine": [18, 19, 137, 138], "minimum": [22, 23], "kth": [24, 115], "partit": 24, "order": [24, 115, 164, 171, 190, 191], "undefin": [24, 115], "sort": [24, 25, 115], "partiton": 24, "flatten": [24, 25, 115, 140, 148, 149, 178], "dimension": [26, 76, 77, 78, 79, 80, 81, 85, 86, 87, 160, 161, 162, 166], "overload": 26, "val": [26, 88], "tupl": [26, 47, 57, 61, 63, 73, 75, 92, 114, 128, 144, 153, 155, 161, 171, 178, 179, 180, 191, 196], "ndarrai": [26, 197], "properti": [27, 35, 44, 49, 51], "argument": [27, 47, 57, 73, 89, 130, 153, 179, 190, 194, 198], "elment": 51, "indices_or_sect": [52, 141], "nest": [56, 190, 191, 196], "correpsond": 56, "ddof": [58, 154], "equal_nan": 59, "nan": 59, "pad": [62, 63, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 160, 161], "dilat": [62, 63], "group": [62, 63, 164], "1d": [62, 64, 149], "convolut": [62, 63, 64, 160, 161], "channel": [62, 63, 160, 161], "c_in": [62, 63], "c_out": [62, 63], "convolv": [62, 63], "2d": 63, "spatial": [63, 164], "symmetr": 63, "discret": [64, 76, 77, 78, 79, 80, 81, 85, 86, 87, 162], "swap": 64, "conv": 64, "filter": [64, 160, 161, 191], "flip": 64, "signal": 64, "divis": 69, "quotient": 69, "mathrm": [71, 135], "frac": [71, 135, 164, 165, 168, 174], "pi": 71, "int_0": 71, "dx": 71, "erf": 72, "retain_graph": [73, 130], "node": [73, 156], "dict": [73, 95, 131, 191, 196], "leaf": [73, 178, 191], "preserv": [73, 128], "intend": 73, "control": [73, 194], "flow": [73, 145], "exponenti": 74, "insert": 75, "One": [76, 79, 85, 129], "fourier": [76, 77, 78, 79, 80, 81, 85, 86, 87], "truncat": [76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 125], "zero": [76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 159, 190], "dft": [76, 77, 78, 79, 80, 81, 85, 86, 87], "rfft": 82, "real": [82, 83, 84, 85, 86, 87], "rfft2": 83, "rfftn": 84, "silent": [85, 86, 87], "fun": [89, 92, 153, 155, 156], "cpp_function": [89, 153, 156], "variabl": [6, 89, 92, 153, 155, 156], "strict": [90, 93, 191], "binari": [95, 130, 131, 132], "npy": [95, 130], "natur": [96, 98], "logarithm": [96, 97, 98, 99], "log": [98, 100, 102], "plu": 98, "exp": [100, 102, 119, 139], "stabl": [100, 102, 139], "multipl": [6, 103, 109, 167], "prepend": 103, "remov": [103, 118, 144], "anoth": [103, 146, 157, 191], "negat": 110, "pad_width": 114, "constant_valu": 114, "edg": 114, "before_1": 114, "after_1": 114, "before_2": 114, "after_2": 114, "before_n": 114, "after_n": 114, "integ": [114, 117, 122, 141, 156, 162, 187], "before_i": 114, "after_i": 114, "extend": 114, "side": 114, "smaller": 115, "distribut": [6, 117, 118, 119, 121, 125, 126], "prng": [117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 194], "num_sampl": 118, "unnorm": 118, "draw": 118, "uint32": [118, 187], "cdf": [119, 163, 181], "accord": [119, 157, 167], "seed": 120, "low": [122, 126], "high": [122, 126, 162, 190], "probabl": [6, 122], "lower": [122, 125, 126], "upper": [122, 125, 126], "bound": [122, 125, 126, 163], "roadcast": 122, "domain": 125, "optino": 125, "uniformli": 126, "reciproc": 129, "arr": 130, "retain": 130, "dure": 130, "uncompress": 131, "my_path": 131, "tree_flatten": [131, 180], "transformerencod": 131, "128": [131, 190], "flat_param": 131, "keyword": [89, 131, 132, 153, 190, 194, 198], "compress": 132, "subarrai": 141, "being": [145, 190], "ident": 145, "prevent": 145, "unchang": [145, 170], "taken": 148, "prior": [148, 149], "equial": 148, "exclud": 149, "mse": 153, "param": [153, 190], "lvalu": 153, "dlvalu": 153, "dparam": 153, "lasso": 153, "l1": 153, "varianc": [154, 164], "divisor": 154, "cotang": 155, "in_ax": 156, "out_ax": 156, "prefix": [156, 178], "select": [157, 191], "in_channel": [160, 161], "out_channel": [160, 161], "kernel_s": [160, 161], "appli": [160, 161, 163, 164, 165, 166, 168, 169, 172, 175, 179, 181, 184, 185, 191], "nlc": 160, "learnabl": [160, 161, 171], "nhwc": 161, "height": 161, "width": 161, "num_embed": 162, "lookup": 162, "tabl": [162, 187], "typic": [162, 193], "usual": [162, 196], "vocabulari": 162, "approx": 163, "unit": [163, 169, 172, 181, 182, 183, 184, 185], "textrm": [163, 181], "phi": [163, 181], "geluapprox": 163, "sigma": [163, 172, 182, 183, 185], "60033": [163, 182], "0433603": [163, 182], "gelufast": 163, "773": [163, 183], "gelu_approx": [163, 181], "gelu_fast_approx": [163, 181], "regard": 163, "num_group": 164, "affin": [164, 165, 166], "pytorch_compat": 164, "var": [164, 165], "epsilon": [164, 165, 168, 174], "gamma": [164, 165, 168], "particular": 164, "split": 164, "preced": 164, "http": [164, 165, 168, 170], "org": [164, 165, 168, 170], "1803": 164, "08494": 164, "stabil": [164, 165, 168], "1607": 165, "06450": 165, "query_input_dim": 167, "key_input_dim": 167, "value_input_dim": 167, "value_dim": 167, "value_output_dim": 167, "dot": [167, 178, 191], "attent": [167, 191], "head": 167, "aggreg": 167, "lineari": 167, "bias": [167, 191], "inf": 167, "neg": 167, "attend": 167, "1910": 168, "07467": 168, "rectifi": [169, 184], "rotat": 170, "consecut": 170, "larger": 170, "slightli": 170, "callabl": [171, 173, 178, 179, 191], "plain": 171, "cdot": [172, 182, 183, 185], "fn": [173, 179, 197], "wrt": 173, "whose": 173, "9": 174, "999": 174, "paper": 174, "omit": 174, "estim": 174, "m_": 174, "beta_1": 174, "m_t": 174, "g_t": [174, 177], "v_": [174, 177], "beta_2": 174, "v_t": [174, 177], "w_": [174, 177], "w_t": [174, 177], "lambda": [174, 177, 179, 191], "kingma": 174, "ba": 174, "2015": 174, "stochast": [174, 177], "iclr": 174, "basi": 175, "optimizerst": 175, "recurs": [176, 190, 191], "defaultdict": 176, "miss": 176, "contrast": 176, "present": 176, "momentum": 177, "descent": 177, "mu": 177, "strength": 177, "is_leaf": 178, "notat": [178, 191], "arbitrari": [178, 191], "depth": 178, "hello": [178, 180], "charact": 178, "flat": [178, 180], "everi": 179, "superset": 179, "extra": 179, "closer": 179, "dict_kei": 179, "recreat": 180, "world": 180, "42": 180, "faster": 181, "gelu": [182, 183], "exact": [182, 183], "0003": 182, "015": 183, "show": [6, 187], "byte": 187, "bool_": 187, "uint8": 187, "unsign": 187, "uint16": 187, "16": [187, 191], "int8": 187, "sign": 187, "int16": 187, "int64": 187, "64": 187, "arm": [6, 187], "arbitrarili": [190, 196, 197], "done": 190, "manual": 190, "explicitli": [190, 194], "solv": 190, "intuit": 190, "freez": [190, 191], "finetun": 190, "in_dim": [190, 191], "out_dim": [190, 191], "enumer": 190, "caus": 190, "local": 190, "scope": 190, "l2_loss": 190, "y_hat": 190, "trainable_paramet": [190, 191], "loss_and_grad": 190, "workhors": 190, "Its": 190, "frozen": [190, 191], "subset": [190, 191], "individu": 190, "action": 190, "preclud": 190, "pure": [190, 193], "pattern": 190, "achiev": 190, "other_input": 190, "necessari": 190, "wrap": 190, "subclass": 191, "concept": 191, "mymlp": 191, "in_proj": 191, "map_fn": 191, "filter_fn": 191, "valid_parameter_filt": 191, "apply_to_modul": 191, "apply_fn": 191, "children": 191, "descend": 191, "filter_and_map": 191, "is_leaf_fn": 191, "content": [6, 191], "found": 191, "whether": 191, "drop": 191, "idempot": 191, "ie": 191, "noop": 191, "unfreez": 191, "endswith": 191, "leaf_modul": 191, "load_weight": 191, "named_modul": 191, "save_weight": 191, "unfrozen": 191, "chang": 191, "tracer": 191, "partial": 191, "subsequ": 193, "implicit": 194, "fine": 194, "grain": 194, "manag": 194, "uniform": 194, "pseudo": 194, "altern": 194, "splittabl": 194, "threefri": 194, "counter": 194, "cycl": 196, "inspect": 197, "composit": 197, "sin": 197, "default_stream": 198, "default_devic": 198, "my_devic": 198, "brought": 5, "research": 5, "maco": 6, "13": 6, "recommend": 6, "14": 6, "sonoma": 6, "xcode": 6, "15": 6, "wish": 6, "environ": 6, "export": 6, "developer_dir": 6, "app": 6, "sdk": 6, "xcrun": 6, "macosx": 6, "meet": 6, "seri": 6, "chip": 6, "nativ": 6, "platform": 6, "processor": 6, "i386": 6, "switch": 6, "argnam": [89, 153], "neither": [89, 153]}, "objects": {"mlx.core": [[7, 0, 1, "", "Device"], [8, 0, 1, "", "Dtype"], [9, 0, 1, "", "Stream"], [10, 2, 1, "", "abs"], [11, 2, 1, "", "add"], [12, 2, 1, "", "all"], [13, 2, 1, "", "allclose"], [14, 2, 1, "", "any"], [15, 2, 1, "", "arange"], [16, 2, 1, "", "arccos"], [17, 2, 1, "", "arccosh"], [18, 2, 1, "", "arcsin"], [19, 2, 1, "", "arcsinh"], [20, 2, 1, "", "arctan"], [21, 2, 1, "", "arctanh"], [22, 2, 1, "", "argmax"], [23, 2, 1, "", "argmin"], [24, 2, 1, "", "argpartition"], [25, 2, 1, "", "argsort"], [26, 0, 1, "", "array"], [59, 2, 1, "", "array_equal"], [60, 2, 1, "", "broadcast_to"], [61, 2, 1, "", "concatenate"], [62, 2, 1, "", "conv1d"], [63, 2, 1, "", "conv2d"], [64, 2, 1, "", "convolve"], [65, 2, 1, "", "cos"], [66, 2, 1, "", "cosh"], [67, 2, 1, "", "default_device"], [68, 2, 1, "", "default_stream"], [69, 2, 1, "", "divide"], [70, 2, 1, "", "equal"], [71, 2, 1, "", "erf"], [72, 2, 1, "", "erfinv"], [73, 2, 1, "", "eval"], [74, 2, 1, "", "exp"], [75, 2, 1, "", "expand_dims"], [88, 2, 1, "", "full"], [89, 2, 1, "", "grad"], [90, 2, 1, "", "greater"], [91, 2, 1, "", "greater_equal"], [92, 2, 1, "", "jvp"], [93, 2, 1, "", "less"], [94, 2, 1, "", "less_equal"], [95, 2, 1, "", "load"], [96, 2, 1, "", "log"], [97, 2, 1, "", "log10"], [98, 2, 1, "", "log1p"], [99, 2, 1, "", "log2"], [100, 2, 1, "", "logaddexp"], [101, 2, 1, "", "logical_not"], [102, 2, 1, "", "logsumexp"], [103, 2, 1, "", "matmul"], [104, 2, 1, "", "max"], [105, 2, 1, "", "maximum"], [106, 2, 1, "", "mean"], [107, 2, 1, "", "min"], [108, 2, 1, "", "minimum"], [109, 2, 1, "", "multiply"], [110, 2, 1, "", "negative"], [111, 2, 1, "", "new_stream"], [112, 2, 1, "", "ones"], [113, 2, 1, "", "ones_like"], [114, 2, 1, "", "pad"], [115, 2, 1, "", "partition"], [116, 2, 1, "", "prod"], [127, 2, 1, "", "reciprocal"], [128, 2, 1, "", "reshape"], [129, 2, 1, "", "rsqrt"], [130, 2, 1, "", "save"], [131, 2, 1, "", "savez"], [132, 2, 1, "", "savez_compressed"], [133, 2, 1, "", "set_default_device"], [134, 2, 1, "", "set_default_stream"], [135, 2, 1, "", "sigmoid"], [136, 2, 1, "", "sign"], [137, 2, 1, "", "sin"], [138, 2, 1, "", "sinh"], [139, 2, 1, "", "softmax"], [140, 2, 1, "", "sort"], [141, 2, 1, "", "split"], [142, 2, 1, "", "sqrt"], [143, 2, 1, "", "square"], [144, 2, 1, "", "squeeze"], [145, 2, 1, "", "stop_gradient"], [146, 2, 1, "", "subtract"], [147, 2, 1, "", "sum"], [148, 2, 1, "", "take"], [149, 2, 1, "", "take_along_axis"], [150, 2, 1, "", "tan"], [151, 2, 1, "", "tanh"], [152, 2, 1, "", "transpose"], [153, 2, 1, "", "value_and_grad"], [154, 2, 1, "", "var"], [155, 2, 1, "", "vjp"], [156, 2, 1, "", "vmap"], [157, 2, 1, "", "where"], [158, 2, 1, "", "zeros"], [159, 2, 1, "", "zeros_like"]], "mlx.core.Device": [[7, 1, 1, "", "__init__"]], "mlx.core.Dtype": [[8, 1, 1, "", "__init__"]], "mlx.core.Stream": [[9, 1, 1, "", "__init__"]], "mlx.core.array": [[27, 3, 1, "", "T"], [26, 1, 1, "", "__init__"], [28, 1, 1, "", "abs"], [29, 1, 1, "", "all"], [30, 1, 1, "", "any"], [31, 1, 1, "", "argmax"], [32, 1, 1, "", "argmin"], [33, 1, 1, "", "astype"], [34, 1, 1, "", "cos"], [35, 3, 1, "", "dtype"], [36, 1, 1, "", "exp"], [37, 1, 1, "", "item"], [38, 1, 1, "", "log"], [39, 1, 1, "", "log1p"], [40, 1, 1, "", "logsumexp"], [41, 1, 1, "", "max"], [42, 1, 1, "", "mean"], [43, 1, 1, "", "min"], [44, 3, 1, "", "ndim"], [45, 1, 1, "", "prod"], [46, 1, 1, "", "reciprocal"], [47, 1, 1, "", "reshape"], [48, 1, 1, "", "rsqrt"], [49, 3, 1, "", "shape"], [50, 1, 1, "", "sin"], [51, 3, 1, "", "size"], [52, 1, 1, "", "split"], [53, 1, 1, "", "sqrt"], [54, 1, 1, "", "square"], [55, 1, 1, "", "sum"], [56, 1, 1, "", "tolist"], [57, 1, 1, "", "transpose"], [58, 1, 1, "", "var"]], "mlx.core.fft": [[76, 2, 1, "", "fft"], [77, 2, 1, "", "fft2"], [78, 2, 1, "", "fftn"], [79, 2, 1, "", "ifft"], [80, 2, 1, "", "ifft2"], [81, 2, 1, "", "ifftn"], [82, 2, 1, "", "irfft"], [83, 2, 1, "", "irfft2"], [84, 2, 1, "", "irfftn"], [85, 2, 1, "", "rfft"], [86, 2, 1, "", "rfft2"], [87, 2, 1, "", "rfftn"]], "mlx.core.random": [[117, 2, 1, "", "bernoulli"], [118, 2, 1, "", "categorical"], [119, 2, 1, "", "gumbel"], [120, 2, 1, "", "key"], [121, 2, 1, "", "normal"], [122, 2, 1, "", "randint"], [123, 2, 1, "", "seed"], [124, 2, 1, "", "split"], [125, 2, 1, "", "truncated_normal"], [126, 2, 1, "", "uniform"]], "mlx.nn": [[160, 0, 1, "", "Conv1d"], [161, 0, 1, "", "Conv2d"], [162, 0, 1, "", "Embedding"], [163, 0, 1, "", "GELU"], [164, 0, 1, "", "GroupNorm"], [165, 0, 1, "", "LayerNorm"], [166, 0, 1, "", "Linear"], [191, 0, 1, "", "Module"], [167, 0, 1, "", "MultiHeadAttention"], [168, 0, 1, "", "RMSNorm"], [169, 0, 1, "", "ReLU"], [170, 0, 1, "", "RoPE"], [171, 0, 1, "", "Sequential"], [172, 0, 1, "", "SiLU"], [181, 0, 1, "", "gelu"], [182, 0, 1, "", "gelu_approx"], [183, 0, 1, "", "gelu_fast_approx"], [184, 0, 1, "", "relu"], [185, 0, 1, "", "silu"], [173, 2, 1, "", "value_and_grad"]], "mlx.nn.Module": [[191, 1, 1, "", "apply"], [191, 1, 1, "", "apply_to_modules"], [191, 1, 1, "", "children"], [191, 1, 1, "", "filter_and_map"], [191, 1, 1, "", "freeze"], [191, 1, 1, "", "leaf_modules"], [191, 1, 1, "", "load_weights"], [191, 1, 1, "", "modules"], [191, 1, 1, "", "named_modules"], [191, 1, 1, "", "parameters"], [191, 1, 1, "", "save_weights"], [191, 1, 1, "", "trainable_parameters"], [191, 1, 1, "", "unfreeze"], [191, 1, 1, "", "update"]], "mlx.optimizers": [[174, 0, 1, "", "Adam"], [175, 0, 1, "", "Optimizer"], [176, 0, 1, "", "OptimizerState"], [177, 0, 1, "", "SGD"]], "mlx.optimizers.Optimizer": [[175, 4, 1, "", "state"]], "mlx.utils": [[178, 2, 1, "", "tree_flatten"], [179, 2, 1, "", "tree_map"], [180, 2, 1, "", "tree_unflatten"]]}, "objtypes": {"0": "py:class", "1": "py:method", "2": "py:function", "3": "py:property", "4": "py:attribute"}, "objnames": {"0": ["py", "class", "Python class"], "1": ["py", "method", "Python method"], "2": ["py", "function", "Python function"], "3": ["py", "property", "Python property"], "4": ["py", "attribute", "Python attribute"]}, "titleterms": {"oper": [0, 1, 192], "develop": 1, "document": 1, "introduc": 1, "exampl": [1, 5], "primit": 1, "us": [1, 198], "implement": [1, 3], "cpu": 1, "backend": 1, "gpu": 1, "transform": [1, 195, 197], "build": [1, 6], "bind": 1, "python": [1, 5, 6], "cmake": 1, "setuptool": 1, "usag": [1, 5], "result": 1, "script": [1, 3], "download": [1, 3], "code": [1, 3], "linear": [2, 166], "regress": 2, "llm": 3, "infer": 3, "model": 3, "attent": 3, "layer": [3, 4, 190], "encod": 3, "full": [3, 88], "gener": 3, "put": 3, "all": [3, 12, 29], "togeth": 3, "convert": 3, "weight": 3, "load": [3, 95], "benchmark": 3, "multi": 4, "perceptron": 4, "mlx": [5, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 127, 128, 129, 130, 131, 132, 133, 134, 135, 136, 137, 138, 139, 140, 141, 142, 143, 144, 145, 146, 147, 148, 149, 150, 151, 152, 153, 154, 155, 156, 157, 158, 159, 160, 161, 162, 163, 164, 165, 166, 167, 168, 169, 170, 171, 172, 173, 174, 175, 176, 177, 178, 179, 180, 181, 182, 183, 184, 185, 191], "instal": [5, 6], "api": [5, 6], "refer": 5, "c": [5, 6], "further": 5, "read": 5, "from": 6, "pypi": 6, "sourc": 6, "requir": 6, "option": 6, "core": [7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 127, 128, 129, 130, 131, 132, 133, 134, 135, 136, 137, 138, 139, 140, 141, 142, 143, 144, 145, 146, 147, 148, 149, 150, 151, 152, 153, 154, 155, 156, 157, 158, 159], "devic": [7, 188, 197], "dtype": [8, 35], "stream": [9, 188, 197, 198], "ab": [10, 28], "add": 11, "allclos": 13, "ani": [14, 30], "arang": 15, "arcco": 16, "arccosh": 17, "arcsin": 18, "arcsinh": 19, "arctan": 20, "arctanh": 21, "argmax": [22, 31], "argmin": [23, 32], "argpartit": 24, "argsort": 25, "arrai": [26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 186], "t": 27, "astyp": 33, "co": [34, 65], "exp": [36, 74], "item": 37, "log": [38, 96], "log1p": [39, 98], "logsumexp": [40, 102], "max": [41, 104], "mean": [42, 106], "min": [43, 107], "ndim": 44, "prod": [45, 116], "reciproc": [46, 127], "reshap": [47, 128], "rsqrt": [48, 129], "shape": 49, "sin": [50, 137], "size": 51, "split": [52, 124, 141], "sqrt": [53, 142], "squar": [54, 143], "sum": [55, 147], "tolist": 56, "transpos": [57, 152], "var": [58, 154], "array_equ": 59, "broadcast_to": 60, "concaten": 61, "conv1d": [62, 160], "conv2d": [63, 161], "convolv": 64, "cosh": 66, "default_devic": 67, "default_stream": 68, "divid": 69, "equal": 70, "erf": 71, "erfinv": 72, "eval": 73, "expand_dim": 75, "fft": [76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 189], "fft2": 77, "fftn": 78, "ifft": 79, "ifft2": 80, "ifftn": 81, "irfft": 82, "irfft2": 83, "irfftn": 84, "rfft": 85, "rfft2": 86, "rfftn": 87, "grad": [89, 190], "greater": 90, "greater_equ": 91, "jvp": 92, "less": 93, "less_equ": 94, "log10": 97, "log2": 99, "logaddexp": 100, "logical_not": 101, "matmul": 103, "maximum": 105, "minimum": 108, "multipli": 109, "neg": 110, "new_stream": 111, "ones": 112, "ones_lik": 113, "pad": 114, "partit": 115, "random": [117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 194], "bernoulli": 117, "categor": 118, "gumbel": 119, "kei": 120, "normal": 121, "randint": 122, "seed": 123, "truncated_norm": 125, "uniform": 126, "save": 130, "savez": 131, "savez_compress": 132, "set_default_devic": 133, "set_default_stream": 134, "sigmoid": 135, "sign": 136, "sinh": 138, "softmax": 139, "sort": 140, "squeez": 144, "stop_gradi": 145, "subtract": 146, "take": 148, "take_along_axi": 149, "tan": 150, "tanh": 151, "value_and_grad": [153, 173], "vjp": 155, "vmap": 156, "where": 157, "zero": 158, "zeros_lik": 159, "nn": [160, 161, 162, 163, 164, 165, 166, 167, 168, 169, 170, 171, 172, 173, 181, 182, 183, 184, 185, 191], "embed": 162, "gelu": [163, 181], "groupnorm": 164, "layernorm": 165, "multiheadattent": 167, "rmsnorm": 168, "relu": [169, 184], "rope": 170, "sequenti": 171, "silu": [172, 185], "optim": [174, 175, 176, 177, 193], "adam": 174, "optimizerst": 176, "sgd": 177, "util": [178, 179, 180, 196], "tree_flatten": 178, "tree_map": 179, "tree_unflatten": 180, "gelu_approx": 182, "gelu_fast_approx": 183, "data": 187, "type": 187, "support": 187, "neural": 190, "network": 190, "quick": [190, 197], "start": [190, 197], "The": 190, "modul": [190, 191], "class": 190, "paramet": 190, "updat": 190, "valu": 190, "tree": 196, "guid": 197, "basic": 197, "function": 197, "graph": 197, "specifi": 198, "troubleshoot": 6}, "envversion": {"sphinx.domains.c": 2, "sphinx.domains.changeset": 1, "sphinx.domains.citation": 1, "sphinx.domains.cpp": 6, "sphinx.domains.index": 1, "sphinx.domains.javascript": 2, "sphinx.domains.math": 2, "sphinx.domains.python": 3, "sphinx.domains.rst": 2, "sphinx.domains.std": 2, "sphinx.ext.intersphinx": 1, "sphinx": 56}})
\ No newline at end of file
diff --git a/docs/build/html/using_streams.html b/docs/build/html/using_streams.html
index 1a9ea3dac..6af3f56d1 100644
--- a/docs/build/html/using_streams.html
+++ b/docs/build/html/using_streams.html
@@ -9,7 +9,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>Using Streams &#8212; MLX 0.0.3 documentation</title>
+    <title>Using Streams &#8212; MLX 0.0.4 documentation</title>
   
   
   
@@ -134,8 +134,8 @@
       
     
     
-    <img src="_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.3 documentation - Home"/>
-    <script>document.write(`<img src="_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.3 documentation - Home"/>`);</script>
+    <img src="_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.0.4 documentation - Home"/>
+    <script>document.write(`<img src="_static/mlx_logo.png" class="logo__image only-dark" alt="MLX 0.0.4 documentation - Home"/>`);</script>
   
   
 </a></div>