mlx 1.0.0

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

@@ -0,0 +1,186 @@
+.. _nn:
+
+.. currentmodule:: mlx.nn
+
+Neural Networks
+===============
+
+Writing arbitrarily complex neural networks in MLX can be done using only
+:class:`mlx.core.array` and :meth:`mlx.core.value_and_grad`. However, this requires the
+user to write again and again the same simple neural network operations as well
+as handle all the parameter state and initialization manually and explicitly.
+
+The module :mod:`mlx.nn` solves this problem by providing an intuitive way of
+composing neural network layers, initializing their parameters, freezing them
+for finetuning and more.
+
+Quick Start with Neural Networks
+---------------------------------
+
+.. code-block:: python
+
+    import mlx.core as mx
+    import mlx.nn as nn
+
+    class MLP(nn.Module):
+        def __init__(self, in_dims: int, out_dims: int):
+            super().__init__()
+
+            self.layers = [
+                nn.Linear(in_dims, 128),
+                nn.Linear(128, 128),
+                nn.Linear(128, out_dims),
+            ]
+
+        def __call__(self, x):
+            for i, l in enumerate(self.layers):
+                x = mx.maximum(x, 0) if i > 0 else x
+                x = l(x)
+            return x
+
+    # The model is created with all its parameters but nothing is initialized
+    # yet because MLX is lazily evaluated
+    mlp = MLP(2, 10)
+
+    # We can access its parameters by calling mlp.parameters()
+    params = mlp.parameters()
+    print(params["layers"][0]["weight"].shape)
+
+    # Printing a parameter will cause it to be evaluated and thus initialized
+    print(params["layers"][0])
+
+    # We can also force evaluate all parameters to initialize the model
+    mx.eval(mlp.parameters())
+
+    # A simple loss function.
+    # NOTE: It doesn't matter how it uses the mlp model. It currently captures
+    #       it from the local scope. It could be a positional argument or a
+    #       keyword argument.
+    def l2_loss(x, y):
+        y_hat = mlp(x)
+        return (y_hat - y).square().mean()
+
+    # Calling `nn.value_and_grad` instead of `mx.value_and_grad` returns the
+    # gradient with respect to `mlp.trainable_parameters()`
+    loss_and_grad = nn.value_and_grad(mlp, l2_loss)
+
+.. _module_class:
+
+The Module Class
+----------------
+
+The workhorse of any neural network library is the :class:`Module` class. In
+MLX the :class:`Module` class is a container of :class:`mlx.core.array` or
+:class:`Module` instances. Its main function is to provide a way to
+recursively **access** and **update** its parameters and those of its
+submodules.
+
+Parameters
+^^^^^^^^^^
+
+A parameter of a module is any public member of type :class:`mlx.core.array` (its
+name should not start with ``_``). It can be arbitrarily nested in other
+:class:`Module` instances or lists and dictionaries.
+
+:meth:`Module.parameters` can be used to extract a nested dictionary with all
+the parameters of a module and its submodules.
+
+A :class:`Module` can also keep track of "frozen" parameters. See the
+:meth:`Module.freeze` method for more details. :meth:`mlx.nn.value_and_grad`
+the gradients returned will be with respect to these trainable parameters.
+
+
+Updating the Parameters
+^^^^^^^^^^^^^^^^^^^^^^^
+
+MLX modules allow accessing and updating individual parameters. However, most
+times we need to update large subsets of a module's parameters. This action is
+performed by :meth:`Module.update`.
+
+
+Inspecting Modules
+^^^^^^^^^^^^^^^^^^
+
+The simplest way to see the model architecture is to print it. Following along with
+the above example, you can print the ``MLP`` with:
+
+.. code-block:: python
+
+  print(mlp)
+
+This will display:
+
+.. code-block:: shell
+
+  MLP(
+    (layers.0): Linear(input_dims=2, output_dims=128, bias=True)
+    (layers.1): Linear(input_dims=128, output_dims=128, bias=True)
+    (layers.2): Linear(input_dims=128, output_dims=10, bias=True)
+  )
+
+To get more detailed information on the arrays in a :class:`Module` you can use
+:func:`mlx.utils.tree_map` on the parameters. For example, to see the shapes of
+all the parameters in a :class:`Module` do:
+
+.. code-block:: python
+
+   from mlx.utils import tree_map
+   shapes = tree_map(lambda p: p.shape, mlp.parameters())
+
+As another example, you can count the number of parameters in a :class:`Module`
+with:
+
+.. code-block:: python
+
+   from mlx.utils import tree_flatten
+   num_params = sum(v.size for _, v in tree_flatten(mlp.parameters()))
+
+
+Value and Grad
+--------------
+
+Using a :class:`Module` does not preclude using MLX's high order function
+transformations (:meth:`mlx.core.value_and_grad`, :meth:`mlx.core.grad`, etc.). However,
+these function transformations assume pure functions, namely the parameters
+should be passed as an argument to the function being transformed.
+
+There is an easy pattern to achieve that with MLX modules
+
+.. code-block:: python
+
+    model = ...
+
+    def f(params, other_inputs):
+        model.update(params)  # <---- Necessary to make the model use the passed parameters
+        return model(other_inputs)
+
+    f(model.trainable_parameters(), mx.zeros((10,)))
+
+However, :meth:`mlx.nn.value_and_grad` provides precisely this pattern and only
+computes the gradients with respect to the trainable parameters of the model.
+
+In detail:
+
+- it wraps the passed function with a function that calls :meth:`Module.update`
+  to make sure the model is using the provided parameters.
+- it calls :meth:`mlx.core.value_and_grad` to transform the function into a function
+  that also computes the gradients with respect to the passed parameters.
+- it wraps the returned function with a function that passes the trainable
+  parameters as the first argument to the function returned by
+  :meth:`mlx.core.value_and_grad`
+
+.. autosummary::
+   :toctree: _autosummary
+
+   value_and_grad
+   quantize
+   average_gradients
+
+.. toctree::
+
+   nn/module
+   nn/layers
+   nn/functions
+   nn/losses
+   nn/init
+   nn/distributed

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

@@ -0,0 +1,184 @@
+.. _ops:
+
+Operations
+==========
+
+.. currentmodule:: mlx.core
+
+.. autosummary::
+  :toctree: _autosummary
+
+   abs
+   add
+   addmm
+   all
+   allclose
+   any
+   arange
+   arccos
+   arccosh
+   arcsin
+   arcsinh
+   arctan
+   arctan2
+   arctanh
+   argmax
+   argmin
+   argpartition
+   argsort
+   array_equal
+   as_strided
+   atleast_1d
+   atleast_2d
+   atleast_3d
+   bitwise_and
+   bitwise_invert
+   bitwise_or
+   bitwise_xor
+   block_masked_mm
+   broadcast_arrays
+   broadcast_to
+   ceil
+   clip
+   concatenate
+   contiguous
+   conj
+   conjugate
+   convolve
+   conv1d
+   conv2d
+   conv3d
+   conv_transpose1d
+   conv_transpose2d
+   conv_transpose3d
+   conv_general
+   cos
+   cosh
+   cummax
+   cummin
+   cumprod
+   cumsum
+   degrees
+   dequantize
+   diag
+   diagonal
+   divide
+   divmod
+   einsum
+   einsum_path
+   equal
+   erf
+   erfinv
+   exp
+   expm1
+   expand_dims
+   eye
+   flatten
+   floor
+   floor_divide
+   full
+   gather_mm
+   gather_qmm
+   greater
+   greater_equal
+   hadamard_transform
+   identity
+   imag
+   inner
+   isfinite
+   isclose
+   isinf
+   isnan
+   isneginf
+   isposinf
+   issubdtype
+   kron
+   left_shift
+   less
+   less_equal
+   linspace
+   load
+   log
+   log2
+   log10
+   log1p
+   logaddexp
+   logcumsumexp
+   logical_not
+   logical_and
+   logical_or
+   logsumexp
+   matmul
+   max
+   maximum
+   mean
+   median
+   meshgrid
+   min
+   minimum
+   moveaxis
+   multiply
+   nan_to_num
+   negative
+   not_equal
+   ones
+   ones_like
+   outer
+   partition
+   pad
+   power
+   prod
+   put_along_axis
+   quantize
+   quantized_matmul
+   radians
+   real
+   reciprocal
+   remainder
+   repeat
+   reshape
+   right_shift
+   roll
+   round
+   rsqrt
+   save
+   savez
+   savez_compressed
+   save_gguf
+   save_safetensors
+   sigmoid
+   sign
+   sin
+   sinh
+   slice
+   slice_update
+   softmax
+   sort
+   split
+   sqrt
+   square
+   squeeze
+   stack
+   std
+   stop_gradient
+   subtract
+   sum
+   swapaxes
+   take
+   take_along_axis
+   tan
+   tanh
+   tensordot
+   tile
+   topk
+   trace
+   transpose
+   tri
+   tril
+   triu
+   unflatten
+   var
+   view
+   where
+   zeros
+   zeros_like

data/mlx/docs/src/python/optimizers/common_optimizers.rst

@@ -0,0 +1,22 @@
+.. _common_optimizers:
+
+Common Optimizers
+=================
+
+.. currentmodule:: mlx.optimizers
+
+.. autosummary::
+   :toctree: _autosummary
+   :template: optimizers-template.rst
+
+   SGD
+   RMSprop
+   Adagrad
+   Adafactor
+   AdaDelta
+   Adam
+   AdamW
+   Adamax
+   Lion
+   MultiOptimizer
+   Muon

data/mlx/docs/src/python/optimizers/optimizer.rst

@@ -0,0 +1,23 @@
+Optimizer
+=========
+
+.. currentmodule:: mlx.optimizers
+
+.. autoclass:: Optimizer 
+
+
+   .. rubric:: Attributes
+
+   .. autosummary::
+      :toctree: _autosummary
+
+      Optimizer.state
+   
+   .. rubric:: Methods
+
+   .. autosummary::
+      :toctree: _autosummary
+   
+      Optimizer.apply_gradients
+      Optimizer.init
+      Optimizer.update

data/mlx/docs/src/python/optimizers/schedulers.rst

@@ -0,0 +1,15 @@
+.. _schedulers:
+
+Schedulers
+==========
+
+.. currentmodule:: mlx.optimizers
+
+.. autosummary::
+   :toctree: _autosummary
+
+   cosine_decay    
+   exponential_decay    
+   join_schedules
+   linear_schedule
+   step_decay    

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

@@ -0,0 +1,78 @@
+.. _optimizers:
+
+.. currentmodule:: mlx.optimizers
+
+Optimizers
+==========
+
+The optimizers in MLX can be used both with :mod:`mlx.nn` but also with pure
+:mod:`mlx.core` functions. A typical example involves calling
+:meth:`Optimizer.update` to update a model's parameters based on the loss
+gradients and subsequently calling :func:`mlx.core.eval` to evaluate both the
+model's parameters and the **optimizer state**.
+
+.. code-block:: python
+
+    # Create a model
+    model = MLP(num_layers, train_images.shape[-1], hidden_dim, num_classes)
+    mx.eval(model.parameters())
+
+    # Create the gradient function and the optimizer
+    loss_and_grad_fn = nn.value_and_grad(model, loss_fn)
+    optimizer = optim.SGD(learning_rate=learning_rate)
+
+    for e in range(num_epochs):
+        for X, y in batch_iterate(batch_size, train_images, train_labels):
+            loss, grads = loss_and_grad_fn(model, X, y)
+
+            # Update the model with the gradients. So far no computation has happened.
+            optimizer.update(model, grads)
+
+            # Compute the new parameters but also the optimizer state.
+            mx.eval(model.parameters(), optimizer.state)
+
+Saving and Loading
+------------------
+
+To serialize an optimizer, save its state. To load an optimizer, load and set
+the saved state. Here's a simple example:
+
+.. code-block:: python
+
+   import mlx.core as mx
+   from mlx.utils import tree_flatten, tree_unflatten
+   import mlx.optimizers as optim
+
+   optimizer = optim.Adam(learning_rate=1e-2)
+
+   # Perform some updates with the optimizer
+   model = {"w" : mx.zeros((5, 5))}
+   grads = {"w" : mx.ones((5, 5))}
+   optimizer.update(model, grads)
+
+   # Save the state
+   state = tree_flatten(optimizer.state, destination={})
+   mx.save_safetensors("optimizer.safetensors", state)
+
+   # Later on, for example when loading from a checkpoint,
+   # recreate the optimizer and load the state
+   optimizer = optim.Adam(learning_rate=1e-2)
+
+   state = tree_unflatten(mx.load("optimizer.safetensors"))
+   optimizer.state = state
+
+Note, not every optimizer configuation parameter is saved in the state. For
+example, for Adam the learning rate is saved but the ``betas`` and ``eps``
+parameters are not. A good rule of thumb is if the parameter can be scheduled
+then it will be included in the optimizer state.
+
+.. toctree::
+
+   optimizers/optimizer
+   optimizers/common_optimizers
+   optimizers/schedulers
+
+.. autosummary::
+   :toctree: _autosummary
+
+   clip_grad_norm

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

@@ -0,0 +1,48 @@
+.. _random:
+
+Random
+======
+
+Random sampling functions in MLX use an implicit global PRNG state by default.
+However, all function take an optional ``key`` keyword argument for when more
+fine-grained control or explicit state management is needed.
+
+For example, you can generate random numbers with:
+
+.. code-block:: python
+
+  for _ in range(3):
+    print(mx.random.uniform())
+
+which will print a sequence of unique pseudo random numbers. Alternatively you
+can explicitly set the key:
+
+.. code-block:: python
+
+  key = mx.random.key(0)
+  for _ in range(3):
+    print(mx.random.uniform(key=key))
+
+which will yield the same pseudo random number at each iteration.
+
+Following `JAX's PRNG design <https://jax.readthedocs.io/en/latest/jep/263-prng.html>`_
+we use a splittable version of Threefry, which is a counter-based PRNG.
+
+.. currentmodule:: mlx.core.random
+
+.. autosummary:: 
+  :toctree: _autosummary
+
+   bernoulli
+   categorical
+   gumbel
+   key
+   normal
+   multivariate_normal
+   randint
+   seed
+   split
+   truncated_normal
+   uniform
+   laplace
+   permutation

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

@@ -0,0 +1,22 @@
+.. _transforms:
+
+Transforms
+==========
+
+.. currentmodule:: mlx.core
+
+.. autosummary::
+  :toctree: _autosummary
+
+   eval
+   async_eval
+   compile
+   checkpoint
+   custom_function
+   disable_compile
+   enable_compile
+   grad
+   value_and_grad
+   jvp
+   vjp
+   vmap

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

@@ -0,0 +1,23 @@
+.. _utils:
+
+Tree Utils
+==========
+
+In MLX we consider a python tree to be an arbitrarily nested collection of
+dictionaries, lists and tuples without cycles. Functions in this module that
+return python trees will be using the default python ``dict``, ``list`` and
+``tuple`` but they can usually process objects that inherit from any of these.
+
+.. note::
+   Dictionaries should have keys that are valid python identifiers.
+
+.. currentmodule:: mlx.utils
+
+.. autosummary:: 
+  :toctree: _autosummary
+
+   tree_flatten
+   tree_unflatten
+   tree_map
+   tree_map_with_path
+   tree_reduce