|
|
Quick Start Guide |
|
|
================= |
|
|
|
|
|
|
|
|
Basics |
|
|
------ |
|
|
|
|
|
.. currentmodule:: mlx.core |
|
|
|
|
|
Import ``mlx.core`` and make an :class:`array`: |
|
|
|
|
|
.. code-block:: python |
|
|
|
|
|
>> import mlx.core as mx |
|
|
>> a = mx.array([1, 2, 3, 4]) |
|
|
>> a.shape |
|
|
[4] |
|
|
>> a.dtype |
|
|
int32 |
|
|
>> b = mx.array([1.0, 2.0, 3.0, 4.0]) |
|
|
>> b.dtype |
|
|
float32 |
|
|
|
|
|
Operations in MLX are lazy. The outputs of MLX operations are not computed |
|
|
until they are needed. To force an array to be evaluated use |
|
|
:func:`eval`. Arrays will automatically be evaluated in a few cases. For |
|
|
example, inspecting a scalar with :meth:`array.item`, printing an array, |
|
|
or converting an array from :class:`array` to :class:`numpy.ndarray` all |
|
|
automatically evaluate the array. |
|
|
|
|
|
.. code-block:: python |
|
|
|
|
|
>> c = a + b # c not yet evaluated |
|
|
>> mx.eval(c) # evaluates c |
|
|
>> c = a + b |
|
|
>> print(c) # Also evaluates c |
|
|
array([2, 4, 6, 8], dtype=float32) |
|
|
>> c = a + b |
|
|
>> import numpy as np |
|
|
>> np.array(c) # Also evaluates c |
|
|
array([2., 4., 6., 8.], dtype=float32) |
|
|
|
|
|
|
|
|
See the page on :ref:`Lazy Evaluation <lazy eval>` for more details. |
|
|
|
|
|
Function and Graph Transformations |
|
|
---------------------------------- |
|
|
|
|
|
MLX has standard function transformations like :func:`grad` and :func:`vmap`. |
|
|
Transformations can be composed arbitrarily. For example |
|
|
``grad(vmap(grad(fn)))`` (or any other composition) is allowed. |
|
|
|
|
|
.. code-block:: python |
|
|
|
|
|
>> x = mx.array(0.0) |
|
|
>> mx.sin(x) |
|
|
array(0, dtype=float32) |
|
|
>> mx.grad(mx.sin)(x) |
|
|
array(1, dtype=float32) |
|
|
>> mx.grad(mx.grad(mx.sin))(x) |
|
|
array(-0, dtype=float32) |
|
|
|
|
|
Other gradient transformations include :func:`vjp` for vector-Jacobian products |
|
|
and :func:`jvp` for Jacobian-vector products. |
|
|
|
|
|
Use :func:`value_and_grad` to efficiently compute both a function's output and |
|
|
gradient with respect to the function's input. |
|
|
|