post nanobind docs fixes and some updates (#889)

* post nanobind docs fixes and some updates

* one more doc nit

* fix for stubs and latex

docs/src/conf.py

@@ -36,9 +36,10 @@ intersphinx_mapping = {
 templates_path = ["_templates"]
 html_static_path = ["_static"]
 source_suffix = ".rst"
-master_doc = "index"
+main_doc = "index"
 highlight_language = "python"
 pygments_style = "sphinx"
+add_module_names = False
 
 # -- Options for HTML output -------------------------------------------------
 
@@ -62,11 +63,19 @@ htmlhelp_basename = "mlx_doc"
 
 
 def setup(app):
-    wrapped = app.registry.documenters["function"].can_document_member
+    from sphinx.util import inspect
 
-    def nanobind_function_patch(member: Any, *args, **kwargs) -> bool:
-        return "nanobind.nb_func" in str(type(member)) or wrapped(
-            member, *args, **kwargs
-        )
+    wrapped_isfunc = inspect.isfunction
 
-    app.registry.documenters["function"].can_document_member = nanobind_function_patch
+    def isfunc(obj):
+        type_name = str(type(obj))
+        if "nanobind.nb_method" in type_name or "nanobind.nb_func" in type_name:
+            return True
+        return wrapped_isfunc(obj)
+
+    inspect.isfunction = isfunc
+
+
+# -- Options for LaTeX output ------------------------------------------------
+
+latex_documents = [(main_doc, "MLX.tex", "MLX Documentation", author, "manual")]

docs/src/index.rst

@@ -62,6 +62,7 @@ are the CPU and GPU.
    python/ops
    python/random
    python/transforms
+   python/fast
    python/fft
    python/linalg
    python/metal

docs/src/python/array.rst

@@ -10,9 +10,12 @@ Array
 
     array
     array.astype
+    array.at
     array.item
     array.tolist
     array.dtype
+    array.itemsize
+    array.nbytes
     array.ndim
     array.shape
     array.size
@@ -23,14 +26,24 @@ Array
     array.argmax
     array.argmin
     array.cos
+    array.cummax
+    array.cummin
+    array.cumprod
+    array.cumsum
+    array.diag
+    array.diagonal
     array.dtype
     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 +53,8 @@ Array
     array.split
     array.sqrt
     array.square
+    array.squeeze
+    array.swapaxes
     array.sum
     array.transpose
     array.T

docs/src/python/data_types.rst

@@ -44,9 +44,15 @@ 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

docs/src/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/src/python/ops.rst

@@ -38,6 +38,10 @@ Operations
    conv_general
    cos
    cosh
+   cummax
+   cummin
+   cumprod
+   cumsum
    dequantize
    diag
    diagonal