mlx 1.0.0

data/mlx/docs/src/python/array.rst

@@ -0,0 +1,65 @@
+.. _array:
+
+Array
+=====
+
+.. currentmodule:: mlx.core
+
+.. autosummary:: 
+   :toctree: _autosummary 
+
+    array
+    array.astype
+    array.at
+    array.item
+    array.tolist
+    array.dtype
+    array.itemsize
+    array.nbytes
+    array.ndim
+    array.shape
+    array.size
+    array.real
+    array.imag
+    array.abs
+    array.all
+    array.any
+    array.argmax
+    array.argmin
+    array.conj
+    array.cos
+    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.logcumsumexp
+    array.logsumexp
+    array.max
+    array.mean
+    array.min
+    array.moveaxis
+    array.prod
+    array.reciprocal
+    array.reshape
+    array.round
+    array.rsqrt
+    array.sin
+    array.split
+    array.sqrt
+    array.square
+    array.squeeze
+    array.std
+    array.sum
+    array.swapaxes
+    array.transpose
+    array.T
+    array.var
+    array.view

data/mlx/docs/src/python/cuda.rst

@@ -0,0 +1,9 @@
+CUDA
+=====
+
+.. currentmodule:: mlx.core.cuda
+
+.. autosummary::
+  :toctree: _autosummary
+
+  is_available

data/mlx/docs/src/python/data_types.rst

@@ -0,0 +1,78 @@
+.. _data_types:
+
+Data Types
+==========
+
+.. currentmodule:: mlx.core
+
+The default floating point type is ``float32`` and the default integer type is
+``int32``. The table below shows supported values for :obj:`Dtype`. 
+
+.. list-table:: Supported Data Types 
+   :widths: 5 3 20
+   :header-rows: 1
+
+   * - Type 
+     - Bytes
+     - Description
+   * - ``bool_``
+     - 1 
+     - Boolean (``True``, ``False``) data type
+   * - ``uint8``
+     - 1 
+     - 8-bit unsigned integer 
+   * - ``uint16``
+     - 2 
+     - 16-bit unsigned integer 
+   * - ``uint32``
+     - 4 
+     - 32-bit unsigned integer 
+   * - ``uint64``
+     - 8 
+     - 64-bit unsigned integer 
+   * - ``int8``
+     - 1 
+     - 8-bit signed integer 
+   * - ``int16``
+     - 2 
+     - 16-bit signed integer 
+   * - ``int32``
+     - 4 
+     - 32-bit signed integer 
+   * - ``int64``
+     - 8 
+     - 64-bit signed integer 
+   * - ``bfloat16``
+     - 2 
+     - 16-bit brain float (e8, m7)
+   * - ``float16``
+     - 2 
+     - 16-bit IEEE float (e5, m10)
+   * - ``float32``
+     - 4 
+     - 32-bit float
+   * - ``float64``
+     - 8
+     - 64-bit double
+   * - ``complex64``
+     - 8 
+     - 64-bit complex float
+
+
+.. note::
+
+    Arrays with type ``float64`` only work with CPU operations. Using
+    ``float64`` arrays on the GPU will result in an exception.
+
+
+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
+   finfo

data/mlx/docs/src/python/devices_and_streams.rst

@@ -0,0 +1,21 @@
+.. _devices_and_streams:
+
+Devices and Streams
+===================
+
+.. currentmodule:: mlx.core
+
+.. autosummary::
+  :toctree: _autosummary
+
+   Device
+   Stream
+   default_device
+   set_default_device
+   default_stream
+   new_stream
+   set_default_stream
+   stream
+   synchronize
+   device_count
+   device_info

data/mlx/docs/src/python/distributed.rst

@@ -0,0 +1,22 @@
+.. _distributed:
+
+.. currentmodule:: mlx.core.distributed
+
+Distributed Communication
+==========================
+
+MLX provides a distributed communication package using MPI. The MPI library is
+loaded at runtime; if MPI is available then distributed communication is also
+made available.
+
+.. autosummary::
+   :toctree: _autosummary
+
+    Group
+    is_available
+    init
+    all_sum
+    all_gather
+    send
+    recv
+    recv_like

data/mlx/docs/src/python/export.rst

@@ -0,0 +1,14 @@
+.. _export:
+
+Export Functions
+================
+
+.. currentmodule:: mlx.core
+
+.. autosummary::
+  :toctree: _autosummary
+
+   export_function
+   import_function
+   exporter
+   export_to_dot

data/mlx/docs/src/python/fast.rst

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

data/mlx/docs/src/python/fft.rst

@@ -0,0 +1,24 @@
+.. _fft:
+
+FFT
+===
+
+.. currentmodule:: mlx.core.fft
+
+.. autosummary:: 
+  :toctree: _autosummary
+
+  fft
+  ifft
+  fft2
+  ifft2
+  fftn
+  ifftn
+  rfft
+  irfft
+  rfft2
+  irfft2
+  rfftn
+  irfftn
+  fftshift
+  ifftshift

data/mlx/docs/src/python/linalg.rst

@@ -0,0 +1,27 @@
+.. _linalg:
+
+Linear Algebra
+==============
+
+.. currentmodule:: mlx.core.linalg
+
+.. autosummary::
+   :toctree: _autosummary
+
+    inv
+    tri_inv
+    norm
+    cholesky
+    cholesky_inv
+    cross
+    qr
+    svd
+    eigvals
+    eig
+    eigvalsh
+    eigh
+    lu
+    lu_factor
+    pinv
+    solve
+    solve_triangular

data/mlx/docs/src/python/memory_management.rst

@@ -0,0 +1,16 @@
+Memory Management
+=================
+
+.. currentmodule:: mlx.core
+
+.. autosummary::
+  :toctree: _autosummary
+
+  get_active_memory
+  get_peak_memory
+  reset_peak_memory
+  get_cache_memory
+  set_memory_limit
+  set_cache_limit
+  set_wired_limit
+  clear_cache

data/mlx/docs/src/python/metal.rst

@@ -0,0 +1,12 @@
+Metal
+=====
+
+.. currentmodule:: mlx.core.metal
+
+.. autosummary::
+  :toctree: _autosummary
+
+  is_available
+  device_info
+  start_capture
+  stop_capture

data/mlx/docs/src/python/nn/distributed.rst

@@ -0,0 +1,30 @@
+.. _nn_distributed:
+
+Distributed
+-----------
+
+Helper Routines
+^^^^^^^^^^^^^^^
+
+The :code:`mlx.nn.layers.distributed` package contains helpful routines to 
+create sharded layers from existing :class:`Modules <mlx.nn.Module>`.
+
+.. currentmodule:: mlx.nn.layers.distributed
+.. autosummary::
+   :toctree: _autosummary
+
+   shard_linear
+   shard_inplace
+
+Layers
+^^^^^^
+
+.. currentmodule:: mlx.nn
+.. autosummary::
+   :toctree: _autosummary
+   :template: nn-module-template.rst
+
+   AllToShardedLinear
+   ShardedToAllLinear
+   QuantizedAllToShardedLinear
+   QuantizedShardedToAllLinear

data/mlx/docs/src/python/nn/functions.rst

@@ -0,0 +1,40 @@
+.. _nn_functions:
+
+.. currentmodule:: mlx.nn
+
+Functions
+---------
+
+Layers without parameters (e.g. activation functions) are also provided as
+simple functions.
+
+.. autosummary::
+   :toctree: _autosummary_functions
+   :template: nn-module-template.rst
+
+   elu
+   celu
+   gelu
+   gelu_approx
+   gelu_fast_approx
+   glu
+   hard_shrink
+   hard_tanh
+   hardswish
+   leaky_relu
+   log_sigmoid
+   log_softmax
+   mish
+   prelu
+   relu
+   relu2
+   relu6
+   selu
+   sigmoid
+   silu
+   softmax
+   softmin
+   softplus
+   softshrink
+   step
+   tanh

data/mlx/docs/src/python/nn/init.rst

@@ -0,0 +1,45 @@
+.. _init:
+
+.. currentmodule:: mlx.nn.init
+
+Initializers
+------------
+
+The ``mlx.nn.init`` package contains commonly used initializers for neural
+network parameters. Initializers return a function which can be applied to any
+input :obj:`mlx.core.array` to produce an initialized output.
+
+For example:
+
+.. code:: python
+
+   import mlx.core as mx
+   import mlx.nn as nn
+
+   init_fn = nn.init.uniform()
+
+   # Produces a [2, 2] uniform matrix
+   param = init_fn(mx.zeros((2, 2)))
+
+To re-initialize all the parameter in an :obj:`mlx.nn.Module` from say a uniform 
+distribution, you can do:
+
+.. code:: python
+  
+   import mlx.nn as nn
+   model = nn.Sequential(nn.Linear(5, 10), nn.ReLU(), nn.Linear(10, 5))
+   init_fn = nn.init.uniform(low=-0.1, high=0.1)
+   model.apply(init_fn)
+   
+
+.. autosummary::
+   :toctree: _autosummary
+
+   constant
+   normal
+   uniform
+   identity
+   glorot_normal
+   glorot_uniform
+   he_normal
+   he_uniform

data/mlx/docs/src/python/nn/layers.rst

@@ -0,0 +1,74 @@
+.. _layers:
+
+.. currentmodule:: mlx.nn
+
+Layers
+------
+
+.. autosummary::
+   :toctree: _autosummary
+   :template: nn-module-template.rst
+
+   ALiBi
+   AllToShardedLinear
+   AvgPool1d
+   AvgPool2d
+   AvgPool3d
+   BatchNorm
+   CELU
+   Conv1d
+   Conv2d
+   Conv3d
+   ConvTranspose1d
+   ConvTranspose2d
+   ConvTranspose3d
+   Dropout
+   Dropout2d
+   Dropout3d
+   Embedding
+   ELU
+   GELU
+   GLU
+   GroupNorm
+   GRU
+   HardShrink
+   HardTanh
+   Hardswish
+   InstanceNorm
+   LayerNorm
+   LeakyReLU
+   Linear
+   LogSigmoid
+   LogSoftmax
+   LSTM
+   MaxPool1d
+   MaxPool2d
+   MaxPool3d
+   Mish
+   MultiHeadAttention
+   PReLU
+   QuantizedAllToShardedLinear
+   QuantizedEmbedding
+   QuantizedLinear
+   QuantizedShardedToAllLinear
+   RMSNorm
+   ReLU
+   ReLU2
+   ReLU6
+   RNN
+   RoPE
+   SELU
+   Sequential
+   ShardedToAllLinear
+   Sigmoid
+   SiLU
+   SinusoidalPositionalEncoding
+   Softmin
+   Softshrink
+   Softsign
+   Softmax
+   Softplus
+   Step
+   Tanh
+   Transformer
+   Upsample

data/mlx/docs/src/python/nn/losses.rst

@@ -0,0 +1,25 @@
+.. _losses:
+
+.. currentmodule:: mlx.nn.losses
+
+Loss Functions
+--------------
+
+.. autosummary::
+   :toctree: _autosummary_functions
+   :template: nn-module-template.rst
+
+   binary_cross_entropy
+   cosine_similarity_loss
+   cross_entropy
+   gaussian_nll_loss
+   hinge_loss
+   huber_loss
+   kl_div_loss
+   l1_loss
+   log_cosh_loss
+   margin_ranking_loss
+   mse_loss
+   nll_loss
+   smooth_l1_loss
+   triplet_loss

data/mlx/docs/src/python/nn/module.rst

@@ -0,0 +1,38 @@
+Module
+======
+
+.. currentmodule:: mlx.nn
+
+.. autoclass:: Module
+
+   .. rubric:: Attributes
+
+   .. autosummary::
+      :toctree: _autosummary
+   
+      Module.training
+      Module.state
+   
+   .. rubric:: Methods
+
+   .. autosummary::
+      :toctree: _autosummary
+   
+      Module.apply
+      Module.apply_to_modules
+      Module.children
+      Module.eval
+      Module.filter_and_map
+      Module.freeze
+      Module.leaf_modules
+      Module.load_weights
+      Module.modules
+      Module.named_modules
+      Module.parameters
+      Module.save_weights
+      Module.set_dtype
+      Module.train
+      Module.trainable_parameters
+      Module.unfreeze
+      Module.update
+      Module.update_modules