rebase

docs/build/html/_sources/dev/extensions.rst

@@ -138,13 +138,13 @@ more concrete:
         * representing the vectorized computation and the axis which
         * corresponds to the output vectorized dimension.
         */
-        virtual std::pair<std::vector<array>, std::vector<int>> vmap(
+        std::pair<std::vector<array>, std::vector<int>> vmap(
             const std::vector<array>& inputs,
             const std::vector<int>& axes) override;
 
-        /** Print the primitive. */
-        void print(std::ostream& os) override {
-            os << "Axpby";
+        /** The name of primitive. */
+        const char* name() const override {
+          return "Axpby";
         }
 
         /** Equivalence check **/

docs/build/html/_sources/install.rst

@@ -23,13 +23,6 @@ To install from PyPI you must meet the following requirements:
     MLX is only available on devices running macOS >= 13.5
     It is highly recommended to use macOS 14 (Sonoma)
 
-
-MLX is also available on conda-forge. To install MLX with conda do:
-
-.. code-block:: shell
-
-   conda install conda-forge::mlx
-
 CUDA
 ^^^^
 
@@ -38,8 +31,16 @@ and SM 7.0 (Volta) and up. To install MLX with CUDA support, run:
 
 .. code-block:: shell
 
-    pip install mlx-cuda
+    pip install "mlx[cuda]"
 
+CPU-only (Linux)
+^^^^^^^^^^^^^^^^
+
+For a CPU-only version of MLX that runs on Linux use:
+
+.. code-block:: shell
+
+    pip install "mlx[cpu]"
 
 Troubleshooting
 ^^^^^^^^^^^^^^^

docs/build/html/_sources/python/optimizers/_autosummary/mlx.optimizers.Muon.rst

@@ -0,0 +1,19 @@
+mlx.optimizers.Muon
+===================
+
+.. currentmodule:: mlx.optimizers
+
+.. autoclass:: Muon
+
+   
+
+   
+   .. rubric:: Methods
+
+   .. autosummary::
+   
+         ~Muon.__init__
+         ~Muon.apply_single
+         ~Muon.init_single
+   
+   

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

@@ -19,3 +19,4 @@ Common Optimizers
    Adamax
    Lion
    MultiOptimizer
+   Muon