docs update

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

@@ -223,7 +223,7 @@ Let's re-implement our operation now in terms of our :class:`Axpby` primitive.
             /* const std::vector<int>& shape = */ out_shape,
             /* Dtype dtype = */ out_dtype,
             /* std::unique_ptr<Primitive> primitive = */
-            std::make_unique<Axpby>(to_stream(s), alpha, beta),
+            std::make_shared<Axpby>(to_stream(s), alpha, beta),
             /* const std::vector<array>& inputs = */ broadcasted_inputs);
     }
 

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

@@ -0,0 +1,52 @@
+Metal Debugger
+==============
+
+Profiling is a key step for performance optimization. You can build MLX with
+the ``MLX_METAL_DEBUG`` option to improve the Metal debugging and optimization
+workflow. The ``MLX_METAL_DEBUG`` debug option:
+
+* Records source during Metal compilation, for later inspection while
+  debugging.
+* Labels Metal objects such as command queues, improving capture readability.
+
+The ``metal::start_capture`` function initiates a capture of all MLX GPU work.
+
+.. code-block:: C++
+
+    int main() {
+        metal::start_capture("/Users/Jane/Developer/MLX.gputrace");
+
+        auto a = arange(10.f, 20.f, 1.f, float32);
+        auto b = arange(30.f, 40.f, 1.f, float32);
+        auto c = add(a, b);
+
+        eval(c);
+
+        metal::stop_capture();
+    }
+
+You can open and replay the GPU trace in Xcode. The ``Dependencies`` view
+has a great overview of all operations. Checkout the `Metal debugger
+documentation`_ for more information.
+
+.. image:: ../_static/metal_debugger/capture.png
+    :class: dark-light
+
+Xcode Workflow
+--------------
+
+You can skip saving to a path by running within Xcode. First, generate an Xcode
+project using CMake.
+
+.. code-block::
+
+    mkdir build && cd build
+    cmake .. -DMLX_METAL_DEBUG=ON -G Xcode
+    open mlx.xcodeproj
+
+Select the ``metal_capture`` example schema and run.
+
+.. image:: ../_static/metal_debugger/schema.png
+    :class: dark-light
+
+.. _`Metal debugger documentation`: https://developer.apple.com/documentation/xcode/metal-debugger

docs/build/html/_sources/index.rst

@@ -58,10 +58,12 @@ are the CPU and GPU.
    :maxdepth: 1
 
    python/array
+   python/data_types
    python/devices_and_streams
    python/ops
    python/random
    python/transforms
+   python/fast
    python/fft
    python/linalg
    python/metal
@@ -80,3 +82,4 @@ are the CPU and GPU.
    :maxdepth: 1
 
    dev/extensions
+   dev/metal_debugger

docs/build/html/_sources/install.rst

@@ -70,16 +70,13 @@ To build and install the MLX python library from source, first, clone MLX from
 
    git clone git@github.com:ml-explore/mlx.git mlx && cd mlx
 
-Make sure that you have `pybind11 <https://pybind11.readthedocs.io/en/stable/index.html>`_
-installed. You can install ``pybind11`` with ``pip``, ``brew`` or ``conda`` as follows:
+Install `nanobind <https://nanobind.readthedocs.io/en/latest/>`_ with:
 
 .. code-block:: shell
 
-    pip install "pybind11[global]"
-    conda install pybind11
-    brew install pybind11
+    pip install git+https://github.com/wjakob/nanobind.git
 
-Then simply build and install it using pip:
+Then simply build and install MLX using pip:
 
 .. code-block:: shell
 
@@ -158,6 +155,8 @@ should point to the path to the built metal library.
      - ON
    * - MLX_BUILD_PYTHON_BINDINGS
      - OFF
+   * - MLX_METAL_DEBUG
+     - OFF
 
 
 .. note::

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

@@ -0,0 +1,22 @@
+mlx.core.DtypeCategory
+======================
+
+.. currentmodule:: mlx.core
+
+.. autoclass:: DtypeCategory
+
+   
+   .. automethod:: __init__
+
+   
+   .. rubric:: Methods
+
+   .. autosummary::
+   
+      ~DtypeCategory.__init__
+   
+   
+
+   
+   
+   

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

@@ -0,0 +1,6 @@
+mlx.core.array.at
+=================
+
+.. currentmodule:: mlx.core
+
+.. autoproperty:: array.at

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

@@ -0,0 +1,6 @@
+mlx.core.array.cummax
+=====================
+
+.. currentmodule:: mlx.core
+
+.. automethod:: array.cummax

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

@@ -0,0 +1,6 @@
+mlx.core.array.cummin
+=====================
+
+.. currentmodule:: mlx.core
+
+.. automethod:: array.cummin

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

@@ -0,0 +1,6 @@
+mlx.core.array.cumprod
+======================
+
+.. currentmodule:: mlx.core
+
+.. automethod:: array.cumprod

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

@@ -0,0 +1,6 @@
+mlx.core.array.cumsum
+=====================
+
+.. currentmodule:: mlx.core
+
+.. automethod:: array.cumsum

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

@@ -0,0 +1,6 @@
+mlx.core.array.diag
+===================
+
+.. currentmodule:: mlx.core
+
+.. automethod:: array.diag

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

@@ -0,0 +1,6 @@
+mlx.core.array.diagonal
+=======================
+
+.. currentmodule:: mlx.core
+
+.. automethod:: array.diagonal

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

@@ -0,0 +1,6 @@
+mlx.core.array.flatten
+======================
+
+.. currentmodule:: mlx.core
+
+.. automethod:: array.flatten

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

@@ -0,0 +1,6 @@
+mlx.core.array.itemsize
+=======================
+
+.. currentmodule:: mlx.core
+
+.. autoproperty:: array.itemsize

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

@@ -0,0 +1,6 @@
+mlx.core.array.log10
+====================
+
+.. currentmodule:: mlx.core
+
+.. automethod:: array.log10

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

@@ -0,0 +1,6 @@
+mlx.core.array.log2
+===================
+
+.. currentmodule:: mlx.core
+
+.. automethod:: array.log2

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

@@ -0,0 +1,6 @@
+mlx.core.array.moveaxis
+=======================
+
+.. currentmodule:: mlx.core
+
+.. automethod:: array.moveaxis

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

@@ -0,0 +1,6 @@
+mlx.core.array.nbytes
+=====================
+
+.. currentmodule:: mlx.core
+
+.. autoproperty:: array.nbytes

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

@@ -0,0 +1,6 @@
+mlx.core.array.squeeze
+======================
+
+.. currentmodule:: mlx.core
+
+.. automethod:: array.squeeze

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

@@ -0,0 +1,6 @@
+mlx.core.array.swapaxes
+=======================
+
+.. currentmodule:: mlx.core
+
+.. automethod:: array.swapaxes

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

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

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

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

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

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

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

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

docs/build/html/_sources/python/_autosummary/mlx.core.fast.layer_norm.rst

@@ -0,0 +1,6 @@
+mlx.core.fast.layer\_norm
+=========================
+
+.. currentmodule:: mlx.core.fast
+
+.. autofunction:: layer_norm

docs/build/html/_sources/python/_autosummary/mlx.core.fast.rms_norm.rst

@@ -0,0 +1,6 @@
+mlx.core.fast.rms\_norm
+=======================
+
+.. currentmodule:: mlx.core.fast
+
+.. autofunction:: rms_norm

docs/build/html/_sources/python/_autosummary/mlx.core.fast.rope.rst

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

docs/build/html/_sources/python/_autosummary/mlx.core.fast.scaled_dot_product_attention.rst

@@ -0,0 +1,6 @@
+mlx.core.fast.scaled\_dot\_product\_attention
+=============================================
+
+.. currentmodule:: mlx.core.fast
+
+.. autofunction:: scaled_dot_product_attention

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

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

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

@@ -10,27 +10,38 @@ Array
 
     array
     array.astype
+    array.at
     array.item
     array.tolist
     array.dtype
+    array.itemsize
+    array.nbytes
     array.ndim
     array.shape
     array.size
-    Dtype
     array.abs
     array.all
     array.any
     array.argmax
     array.argmin
     array.cos
-    array.dtype
+    array.cummax
+    array.cummin
+    array.cumprod
+    array.cumsum
+    array.diag
+    array.diagonal
     array.exp
+    array.flatten
     array.log
+    array.log10
     array.log1p
+    array.log2
     array.logsumexp
     array.max
     array.mean
     array.min
+    array.moveaxis
     array.prod
     array.reciprocal
     array.reshape
@@ -40,6 +51,8 @@ Array
     array.split
     array.sqrt
     array.square
+    array.squeeze
+    array.swapaxes
     array.sum
     array.transpose
     array.T

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

@@ -1,7 +1,5 @@
 .. _data_types:
 
-:orphan:
-
 Data Types
 ==========
 
@@ -44,9 +42,27 @@ The default floating point type is ``float32`` and the default integer type is
    * - ``int64``
      - 8 
      - 64-bit signed integer 
+   * - ``bfloat16``
+     - 2 
+     - 16-bit brain float (e8, m7)
    * - ``float16``
      - 2 
-     - 16-bit float, only available with `ARM C language extensions <https://developer.arm.com/documentation/101028/0012/3--C-language-extensions?lang=en>`_
+     - 16-bit IEEE float (e5, m10)
    * - ``float32``
      - 4 
      - 32-bit float
+   * - ``complex64``
+     - 8 
+     - 64-bit complex float
+
+
+Data type are aranged in a hierarchy. See the :obj:`DtypeCategory` object
+documentation for more information. Use :func:`issubdtype` to determine if one
+``dtype`` (or category) is a subtype of another category.
+
+.. autosummary::
+   :toctree: _autosummary
+
+   Dtype
+   DtypeCategory
+   issubdtype

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

@@ -0,0 +1,14 @@
+.. _fast:
+
+Fast
+====
+
+.. currentmodule:: mlx.core.fast
+
+.. autosummary:: 
+  :toctree: _autosummary
+
+  rms_norm
+  layer_norm
+  rope
+  scaled_dot_product_attention

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

@@ -0,0 +1,6 @@
+mlx.nn.Module.set\_dtype
+========================
+
+.. currentmodule:: mlx.nn
+
+.. automethod:: Module.set_dtype

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

@@ -30,6 +30,7 @@ Module
       Module.named_modules
       Module.parameters
       Module.save_weights
+      Module.set_dtype
       Module.train
       Module.trainable_parameters
       Module.unfreeze

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

@@ -38,6 +38,10 @@ Operations
    conv_general
    cos
    cosh
+   cummax
+   cummin
+   cumprod
+   cumsum
    dequantize
    diag
    diagonal
@@ -58,10 +62,10 @@ Operations
    identity
    inner
    isclose
-   isnan
-   isposinf
-   isneginf
    isinf
+   isnan
+   isneginf
+   isposinf
    less
    less_equal
    linspace

docs/build/html/_sources/usage/saving_and_loading.rst

@@ -49,7 +49,7 @@ it will be added. You can load the array with:
 
 .. code-block:: shell
 
-   >>> mx.load("array.npy", a)
+   >>> mx.load("array.npy")
    array([1], dtype=float32)
 
 Here's an example of saving several arrays to a single file: