docs update

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