# Table of Contents
- [MLX — MLX 0.24.0 documentation](#mlx-mlx-0-24-0-documentation)
- [Quick Start Guide — MLX 0.24.0 documentation](#quick-start-guide-mlx-0-24-0-documentation)
- [Unified Memory — MLX 0.24.0 documentation](#unified-memory-mlx-0-24-0-documentation)
- [Saving and Loading Arrays — MLX 0.24.0 documentation](#saving-and-loading-arrays-mlx-0-24-0-documentation)
- [Build and Install — MLX 0.24.0 documentation](#build-and-install-mlx-0-24-0-documentation)
- [Lazy Evaluation — MLX 0.24.0 documentation](#lazy-evaluation-mlx-0-24-0-documentation)
- [Indexing Arrays — MLX 0.24.0 documentation](#indexing-arrays-mlx-0-24-0-documentation)
- [Function Transforms — MLX 0.24.0 documentation](#function-transforms-mlx-0-24-0-documentation)
- [Using Streams — MLX 0.24.0 documentation](#using-streams-mlx-0-24-0-documentation)
- [Compilation — MLX 0.24.0 documentation](#compilation-mlx-0-24-0-documentation)
- [Exporting Functions — MLX 0.24.0 documentation](#exporting-functions-mlx-0-24-0-documentation)
- [Conversion to NumPy and Other Frameworks — MLX 0.24.0 documentation](#conversion-to-numpy-and-other-frameworks-mlx-0-24-0-documentation)
- [Linear Regression — MLX 0.24.0 documentation](#linear-regression-mlx-0-24-0-documentation)
- [Multi-Layer Perceptron — MLX 0.24.0 documentation](#multi-layer-perceptron-mlx-0-24-0-documentation)
- [LLM inference — MLX 0.24.0 documentation](#llm-inference-mlx-0-24-0-documentation)
- [Distributed Communication — MLX 0.24.0 documentation](#distributed-communication-mlx-0-24-0-documentation)
- [Array — MLX 0.24.0 documentation](#array-mlx-0-24-0-documentation)
- [Devices and Streams — MLX 0.24.0 documentation](#devices-and-streams-mlx-0-24-0-documentation)
- [Data Types — MLX 0.24.0 documentation](#data-types-mlx-0-24-0-documentation)
- [Export Functions — MLX 0.24.0 documentation](#export-functions-mlx-0-24-0-documentation)
- [Random — MLX 0.24.0 documentation](#random-mlx-0-24-0-documentation)
- [Operations — MLX 0.24.0 documentation](#operations-mlx-0-24-0-documentation)
- [Fast — MLX 0.24.0 documentation](#fast-mlx-0-24-0-documentation)
- [Transforms — MLX 0.24.0 documentation](#transforms-mlx-0-24-0-documentation)
- [FFT — MLX 0.24.0 documentation](#fft-mlx-0-24-0-documentation)
- [Linear Algebra — MLX 0.24.0 documentation](#linear-algebra-mlx-0-24-0-documentation)
- [Metal — MLX 0.24.0 documentation](#metal-mlx-0-24-0-documentation)
- [Distributed Communication — MLX 0.24.0 documentation](#distributed-communication-mlx-0-24-0-documentation)
- [Neural Networks — MLX 0.24.0 documentation](#neural-networks-mlx-0-24-0-documentation)
- [Optimizers — MLX 0.24.0 documentation](#optimizers-mlx-0-24-0-documentation)
- [Tree Utils — MLX 0.24.0 documentation](#tree-utils-mlx-0-24-0-documentation)
- [Metal Debugger — MLX 0.24.0 documentation](#metal-debugger-mlx-0-24-0-documentation)
- [Custom Metal Kernels — MLX 0.24.0 documentation](#custom-metal-kernels-mlx-0-24-0-documentation)
- [Operations — MLX 0.24.0 documentation](#operations-mlx-0-24-0-documentation)
- [Custom Extensions in MLX — MLX 0.24.0 documentation](#custom-extensions-in-mlx-mlx-0-24-0-documentation)
- [Using MLX in C++ — MLX 0.24.0 documentation](#using-mlx-in-c-mlx-0-24-0-documentation)
- [Unknown](#unknown)
- [Unknown](#unknown)
- [Launching Distributed Programs — MLX 0.24.0 documentation](#launching-distributed-programs-mlx-0-24-0-documentation)
---
# MLX — MLX 0.24.0 documentation
[Skip to main content](#main-content)
Back to top
Ctrl+K
[](#)
[](https://github.com/ml-explore/mlx "Source repository")
* [.rst](_sources/index.rst "Download source file")
* .pdf
MLX
===
MLX[#](#mlx "Link to this heading")
====================================
MLX is a NumPy-like array framework designed for efficient and flexible machine learning on Apple silicon, brought to you by Apple machine learning research.
The Python API closely follows NumPy with a few exceptions. MLX also has a fully featured C++ API which closely follows the Python API.
The main differences between MLX and NumPy are:
> * **Composable function transformations**: MLX has composable function transformations for automatic differentiation, automatic vectorization, and computation graph optimization.
>
> * **Lazy computation**: Computations in MLX are lazy. Arrays are only materialized when needed.
>
> * **Multi-device**: Operations can run on any of the supported devices (CPU, GPU, …)
>
The design of MLX is inspired by frameworks like [PyTorch](https://pytorch.org/)
, [Jax](https://github.com/google/jax)
, and [ArrayFire](https://arrayfire.org/)
. A notable difference from these frameworks and MLX is the _unified memory model_. Arrays in MLX live in shared memory. Operations on MLX arrays can be performed on any of the supported device types without performing data copies. Currently supported device types are the CPU and GPU.
Install
* [Build and Install](install.html)
Usage
* [Quick Start Guide](usage/quick_start.html)
* [Lazy Evaluation](usage/lazy_evaluation.html)
* [Unified Memory](usage/unified_memory.html)
* [Indexing Arrays](usage/indexing.html)
* [Saving and Loading Arrays](usage/saving_and_loading.html)
* [Function Transforms](usage/function_transforms.html)
* [Compilation](usage/compile.html)
* [Conversion to NumPy and Other Frameworks](usage/numpy.html)
* [Distributed Communication](usage/distributed.html)
* [Using Streams](usage/using_streams.html)
* [Exporting Functions](usage/export.html)
Examples
* [Linear Regression](examples/linear_regression.html)
* [Multi-Layer Perceptron](examples/mlp.html)
* [LLM inference](examples/llama-inference.html)
Python API Reference
* [Array](python/array.html)
* [Data Types](python/data_types.html)
* [Devices and Streams](python/devices_and_streams.html)
* [Export Functions](python/export.html)
* [Operations](python/ops.html)
* [Random](python/random.html)
* [Transforms](python/transforms.html)
* [Fast](python/fast.html)
* [FFT](python/fft.html)
* [Linear Algebra](python/linalg.html)
* [Metal](python/metal.html)
* [Neural Networks](python/nn.html)
* [Optimizers](python/optimizers.html)
* [Distributed Communication](python/distributed.html)
* [Tree Utils](python/tree_utils.html)
C++ API Reference
* [Operations](cpp/ops.html)
Further Reading
* [Custom Extensions in MLX](dev/extensions.html)
* [Metal Debugger](dev/metal_debugger.html)
* [Custom Metal Kernels](dev/custom_metal_kernels.html)
* [Using MLX in C++](dev/mlx_in_cpp.html)
---
# Quick Start Guide — MLX 0.24.0 documentation
[Skip to main content](#main-content)
Back to top
Ctrl+K
[](../index.html)
[](https://github.com/ml-explore/mlx "Source repository")
* [.rst](../_sources/usage/quick_start.rst "Download source file")
* .pdf
Quick Start Guide
=================
Contents
--------
Quick Start Guide[#](#quick-start-guide "Link to this heading")
================================================================
Basics[#](#basics "Link to this heading")
------------------------------------------
Import `mlx.core` and make an [`array`](../python/_autosummary/mlx.core.array.html#mlx.core.array "mlx.core.array")
:
\>> 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 [`eval()`](../python/_autosummary/mlx.core.eval.html#mlx.core.eval "mlx.core.eval")
. Arrays will automatically be evaluated in a few cases. For example, inspecting a scalar with [`array.item()`](../python/_autosummary/mlx.core.array.item.html#mlx.core.array.item "mlx.core.array.item")
, printing an array, or converting an array from [`array`](../python/_autosummary/mlx.core.array.html#mlx.core.array "mlx.core.array")
to [`numpy.ndarray`](https://numpy.org/doc/stable/reference/generated/numpy.ndarray.html#numpy.ndarray "(in NumPy v2.2)")
all automatically evaluate the array.
\>> 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 [Lazy Evaluation](lazy_evaluation.html#lazy-eval)
for more details.
Function and Graph Transformations[#](#function-and-graph-transformations "Link to this heading")
--------------------------------------------------------------------------------------------------
MLX has standard function transformations like [`grad()`](../python/_autosummary/mlx.core.grad.html#mlx.core.grad "mlx.core.grad")
and [`vmap()`](../python/_autosummary/mlx.core.vmap.html#mlx.core.vmap "mlx.core.vmap")
. Transformations can be composed arbitrarily. For example `grad(vmap(grad(fn)))` (or any other composition) is allowed.
\>> 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 [`vjp()`](../python/_autosummary/mlx.core.vjp.html#mlx.core.vjp "mlx.core.vjp")
for vector-Jacobian products and [`jvp()`](../python/_autosummary/mlx.core.jvp.html#mlx.core.jvp "mlx.core.jvp")
for Jacobian-vector products.
Use [`value_and_grad()`](../python/_autosummary/mlx.core.value_and_grad.html#mlx.core.value_and_grad "mlx.core.value_and_grad")
to efficiently compute both a function’s output and gradient with respect to the function’s input.
Contents
---
# Unified Memory — MLX 0.24.0 documentation
[Skip to main content](#main-content)
Back to top
Ctrl+K
[](../index.html)
[](https://github.com/ml-explore/mlx "Source repository")
* [.rst](../_sources/usage/unified_memory.rst "Download source file")
* .pdf
Unified Memory
==============
Contents
--------
Unified Memory[#](#unified-memory "Link to this heading")
==========================================================
Apple silicon has a unified memory architecture. The CPU and GPU have direct access to the same memory pool. MLX is designed to take advantage of that.
Concretely, when you make an array in MLX you don’t have to specify its location:
a \= mx.random.normal((100,))
b \= mx.random.normal((100,))
Both `a` and `b` live in unified memory.
In MLX, rather than moving arrays to devices, you specify the device when you run the operation. Any device can perform any operation on `a` and `b` without needing to move them from one memory location to another. For example:
mx.add(a, b, stream\=mx.cpu)
mx.add(a, b, stream\=mx.gpu)
In the above, both the CPU and the GPU will perform the same add operation. The operations can (and likely will) be run in parallel since there are no dependencies between them. See [Using Streams](using_streams.html#using-streams)
for more information the semantics of streams in MLX.
In the above `add` example, there are no dependencies between operations, so there is no possibility for race conditions. If there are dependencies, the MLX scheduler will automatically manage them. For example:
c \= mx.add(a, b, stream\=mx.cpu)
d \= mx.add(a, c, stream\=mx.gpu)
In the above case, the second `add` runs on the GPU but it depends on the output of the first `add` which is running on the CPU. MLX will automatically insert a dependency between the two streams so that the second `add` only starts executing after the first is complete and `c` is available.
A Simple Example[#](#a-simple-example "Link to this heading")
--------------------------------------------------------------
Here is a more interesting (albeit slightly contrived example) of how unified memory can be helpful. Suppose we have the following computation:
def fun(a, b, d1, d2):
x \= mx.matmul(a, b, stream\=d1)
for \_ in range(500):
b \= mx.exp(b, stream\=d2)
return x, b
which we want to run with the following arguments:
a \= mx.random.uniform(shape\=(4096, 512))
b \= mx.random.uniform(shape\=(512, 4))
The first `matmul` operation is a good fit for the GPU since it’s more compute dense. The second sequence of operations are a better fit for the CPU, since they are very small and would probably be overhead bound on the GPU.
If we time the computation fully on the GPU, we get 2.8 milliseconds. But if we run the computation with `d1=mx.gpu` and `d2=mx.cpu`, then the time is only about 1.4 milliseconds, about twice as fast. These times were measured on an M1 Max.
Contents
---
# Saving and Loading Arrays — MLX 0.24.0 documentation
[Skip to main content](#main-content)
Back to top
Ctrl+K
[](../index.html)
[](https://github.com/ml-explore/mlx "Source repository")
* [.rst](../_sources/usage/saving_and_loading.rst "Download source file")
* .pdf
Saving and Loading Arrays
=========================
Saving and Loading Arrays[#](#saving-and-loading-arrays "Link to this heading")
================================================================================
MLX supports multiple array serialization formats.
| | | | |
| --- | --- | --- | --- |Serialization Formats[#](#id1 "Link to this table")
| Format | Extension | Function | Notes |
| --- | --- | --- | --- |
| NumPy | `.npy` | [`save()`](../python/_autosummary/mlx.core.save.html#mlx.core.save "mlx.core.save") | Single arrays only |
| NumPy archive | `.npz` | [`savez()`](../python/_autosummary/mlx.core.savez.html#mlx.core.savez "mlx.core.savez")
and [`savez_compressed()`](../python/_autosummary/mlx.core.savez_compressed.html#mlx.core.savez_compressed "mlx.core.savez_compressed") | Multiple arrays |
| Safetensors | `.safetensors` | [`save_safetensors()`](../python/_autosummary/mlx.core.save_safetensors.html#mlx.core.save_safetensors "mlx.core.save_safetensors") | Multiple arrays |
| GGUF | `.gguf` | [`save_gguf()`](../python/_autosummary/mlx.core.save_gguf.html#mlx.core.save_gguf "mlx.core.save_gguf") | Multiple arrays |
The [`load()`](../python/_autosummary/mlx.core.load.html#mlx.core.load "mlx.core.load")
function will load any of the supported serialization formats. It determines the format from the extensions. The output of [`load()`](../python/_autosummary/mlx.core.load.html#mlx.core.load "mlx.core.load")
depends on the format.
Here’s an example of saving a single array to a file:
\>>> a \= mx.array(\[1.0\])
>>> mx.save("array", a)
The array `a` will be saved in the file `array.npy` (notice the extension is automatically added). Including the extension is optional; if it is missing it will be added. You can load the array with:
\>>> mx.load("array.npy")
array(\[1\], dtype\=float32)
Here’s an example of saving several arrays to a single file:
\>>> a \= mx.array(\[1.0\])
>>> b \= mx.array(\[2.0\])
>>> mx.savez("arrays", a, b\=b)
For compatibility with [`numpy.savez()`](https://numpy.org/doc/stable/reference/generated/numpy.savez.html#numpy.savez "(in NumPy v2.2)")
the MLX [`savez()`](../python/_autosummary/mlx.core.savez.html#mlx.core.savez "mlx.core.savez")
takes arrays as arguments. If the keywords are missing, then default names will be provided. This can be loaded with:
\>>> mx.load("arrays.npz")
{'b': array(\[2\], dtype\=float32), 'arr\_0': array(\[1\], dtype\=float32)}
In this case [`load()`](../python/_autosummary/mlx.core.load.html#mlx.core.load "mlx.core.load")
returns a dictionary of names to arrays.
The functions [`save_safetensors()`](../python/_autosummary/mlx.core.save_safetensors.html#mlx.core.save_safetensors "mlx.core.save_safetensors")
and [`save_gguf()`](../python/_autosummary/mlx.core.save_gguf.html#mlx.core.save_gguf "mlx.core.save_gguf")
are similar to [`savez()`](../python/_autosummary/mlx.core.savez.html#mlx.core.savez "mlx.core.savez")
, but they take as input a [`dict`](https://docs.python.org/3/library/stdtypes.html#dict "(in Python v3.13)")
of string names to arrays:
\>>> a \= mx.array(\[1.0\])
>>> b \= mx.array(\[2.0\])
>>> mx.save\_safetensors("arrays", {"a": a, "b": b})
---
# Build and Install — MLX 0.24.0 documentation
[Skip to main content](#main-content)
Back to top
Ctrl+K
[](index.html)
[](https://github.com/ml-explore/mlx "Source repository")
* [.rst](_sources/install.rst "Download source file")
* .pdf
Build and Install
=================
Contents
--------
Build and Install[#](#build-and-install "Link to this heading")
================================================================
Python Installation[#](#python-installation "Link to this heading")
--------------------------------------------------------------------
MLX is available on PyPI. All you have to do to use MLX with your own Apple silicon computer is
pip install mlx
To install from PyPI you must meet the following requirements:
* Using an M series chip (Apple silicon)
* Using a native Python >= 3.9
* macOS >= 13.5
Note
MLX is only available on devices running macOS >= 13.5 It is highly recommended to use macOS 14 (Sonoma)
MLX is also available on conda-forge. To install MLX with conda do:
conda install conda-forge::mlx
### Troubleshooting[#](#troubleshooting "Link to this heading")
_My OS and Python versions are in the required range but pip still does not find a matching distribution._
Probably you are using a non-native Python. The output of
python \-c "import platform; print(platform.processor())"
should be `arm`. If it is `i386` (and you have M series machine) then you are using a non-native Python. Switch your Python to a native Python. A good way to do this is with [Conda](https://stackoverflow.com/q/65415996)
.
Build from source[#](#build-from-source "Link to this heading")
----------------------------------------------------------------
### Build Requirements[#](#build-requirements "Link to this heading")
* A C++ compiler with C++17 support (e.g. Clang >= 5.0)
* [cmake](https://cmake.org/)
– version 3.25 or later, and `make`
* Xcode >= 15.0 and macOS SDK >= 14.0
Note
Ensure your shell environment is native `arm`, not `x86` via Rosetta. If the output of `uname -p` is `x86`, see the [troubleshooting section](#build-shell)
below.
### Python API[#](#python-api "Link to this heading")
To build and install the MLX python library from source, first, clone MLX from [its GitHub repo](https://github.com/ml-explore/mlx)
:
git clone git@github.com:ml-explore/mlx.git mlx && cd mlx
Then simply build and install MLX using pip:
CMAKE\_BUILD\_PARALLEL\_LEVEL\=8 pip install .
For developing, install the package with development dependencies, and use an editable install:
CMAKE\_BUILD\_PARALLEL\_LEVEL\=8 pip install \-e ".\[dev\]"
Once the development dependencies are installed, you can build faster with:
CMAKE\_BUILD\_PARALLEL\_LEVEL\=8 python setup.py build\_ext \--inplace
Run the tests with:
python \-m unittest discover python/tests
Optional: Install stubs to enable auto completions and type checking from your IDE:
python setup.py generate\_stubs
### C++ API[#](#c-api "Link to this heading")
Currently, MLX must be built and installed from source.
Similarly to the python library, to build and install the MLX C++ library start by cloning MLX from [its GitHub repo](https://github.com/ml-explore/mlx)
:
git clone git@github.com:ml-explore/mlx.git mlx && cd mlx
Create a build directory and run CMake and make:
mkdir \-p build && cd build
cmake .. && make \-j
Run tests with:
make test
Install with:
make install
Note that the built `mlx.metallib` file should be either at the same directory as the executable statically linked to `libmlx.a` or the preprocessor constant `METAL_PATH` should be defined at build time and it should point to the path to the built metal library.
| | |
| --- | --- |Build Options[#](#id4 "Link to this table")
| Option | Default |
| --- | --- |
| MLX\_BUILD\_TESTS | ON |
| MLX\_BUILD\_EXAMPLES | OFF |
| MLX\_BUILD\_BENCHMARKS | OFF |
| MLX\_BUILD\_METAL | ON |
| MLX\_BUILD\_CPU | ON |
| MLX\_BUILD\_PYTHON\_BINDINGS | OFF |
| MLX\_METAL\_DEBUG | OFF |
| MLX\_BUILD\_SAFETENSORS | ON |
| MLX\_BUILD\_GGUF | ON |
| MLX\_METAL\_JIT | OFF |
Note
If you have multiple Xcode installations and wish to use a specific one while building, you can do so by adding the following environment variable before building
export DEVELOPER\_DIR\="/path/to/Xcode.app/Contents/Developer/"
Further, you can use the following command to find out which macOS SDK will be used
xcrun \-sdk macosx \--show-sdk-version
#### Binary Size Minimization[#](#binary-size-minimization "Link to this heading")
To produce a smaller binary use the CMake flags `CMAKE_BUILD_TYPE=MinSizeRel` and `BUILD_SHARED_LIBS=ON`.
The MLX CMake build has several additional options to make smaller binaries. For example, if you don’t need the CPU backend or support for safetensors and GGUF, you can do:
cmake .. \\
\-DCMAKE\_BUILD\_TYPE\=MinSizeRel \\
\-DBUILD\_SHARED\_LIBS\=ON \\
\-DMLX\_BUILD\_CPU\=OFF \\
\-DMLX\_BUILD\_SAFETENSORS\=OFF \\
\-DMLX\_BUILD\_GGUF\=OFF \\
\-DMLX\_METAL\_JIT\=ON
THE `MLX_METAL_JIT` flag minimizes the size of the MLX Metal library which contains pre-built GPU kernels. This substantially reduces the size of the Metal library by run-time compiling kernels the first time they are used in MLX on a given machine. Note run-time compilation incurs a cold-start cost which can be anwywhere from a few hundred millisecond to a few seconds depending on the application. Once a kernel is compiled, it will be cached by the system. The Metal kernel cache persists across reboots.
### Troubleshooting[#](#id3 "Link to this heading")
#### Metal not found[#](#metal-not-found "Link to this heading")
You see the following error when you try to build:
error: unable to find utility "metal", not a developer tool or in PATH
To fix this, first make sure you have Xcode installed:
xcode-select \--install
Then set the active developer directory:
sudo xcode-select \--switch /Applications/Xcode.app/Contents/Developer
#### x86 Shell[#](#x86-shell "Link to this heading")
If the output of `uname -p` is `x86` then your shell is running as x86 via Rosetta instead of natively.
To fix this, find the application in Finder (`/Applications` for iTerm, `/Applications/Utilities` for Terminal), right-click, and click “Get Info”. Uncheck “Open using Rosetta”, close the “Get Info” window, and restart your terminal.
Verify the terminal is now running natively the following command:
$ uname \-p
arm
Also check that cmake is using the correct architecture:
$ cmake \--system-information | grep CMAKE\_HOST\_SYSTEM\_PROCESSOR
CMAKE\_HOST\_SYSTEM\_PROCESSOR "arm64"
If you see `"x86_64"`, try re-installing `cmake`. If you see `"arm64"` but the build errors out with “Building for x86\_64 on macOS is not supported.” wipe your build cache with `rm -rf build/` and try again.
Contents
---
# Lazy Evaluation — MLX 0.24.0 documentation
[Skip to main content](#main-content)
Back to top
Ctrl+K
[](../index.html)
[](https://github.com/ml-explore/mlx "Source repository")
* [.rst](../_sources/usage/lazy_evaluation.rst "Download source file")
* .pdf
Lazy Evaluation
===============
Contents
--------
Lazy Evaluation[#](#lazy-evaluation "Link to this heading")
============================================================
Why Lazy Evaluation[#](#why-lazy-evaluation "Link to this heading")
--------------------------------------------------------------------
When you perform operations in MLX, no computation actually happens. Instead a compute graph is recorded. The actual computation only happens if an [`eval()`](../python/_autosummary/mlx.core.eval.html#mlx.core.eval "mlx.core.eval")
is performed.
MLX uses lazy evaluation because it has some nice features, some of which we describe below.
### Transforming Compute Graphs[#](#transforming-compute-graphs "Link to this heading")
Lazy evaluation lets us record a compute graph without actually doing any computations. This is useful for function transformations like [`grad()`](../python/_autosummary/mlx.core.grad.html#mlx.core.grad "mlx.core.grad")
and [`vmap()`](../python/_autosummary/mlx.core.vmap.html#mlx.core.vmap "mlx.core.vmap")
and graph optimizations.
Currently, MLX does not compile and rerun compute graphs. They are all generated dynamically. However, lazy evaluation makes it much easier to integrate compilation for future performance enhancements.
### Only Compute What You Use[#](#only-compute-what-you-use "Link to this heading")
In MLX you do not need to worry as much about computing outputs that are never used. For example:
def fun(x):
a \= fun1(x)
b \= expensive\_fun(a)
return a, b
y, \_ \= fun(x)
Here, we never actually compute the output of `expensive_fun`. Use this pattern with care though, as the graph of `expensive_fun` is still built, and that has some cost associated to it.
Similarly, lazy evaluation can be beneficial for saving memory while keeping code simple. Say you have a very large model `Model` derived from [`mlx.nn.Module`](../python/nn/module.html#mlx.nn.Module "mlx.nn.Module")
. You can instantiate this model with `model = Model()`. Typically, this will initialize all of the weights as `float32`, but the initialization does not actually compute anything until you perform an [`eval()`](../python/_autosummary/mlx.core.eval.html#mlx.core.eval "mlx.core.eval")
. If you update the model with `float16` weights, your maximum consumed memory will be half that required if eager computation was used instead.
This pattern is simple to do in MLX thanks to lazy computation:
model \= Model() \# no memory used yet
model.load\_weights("weights\_fp16.safetensors")
When to Evaluate[#](#when-to-evaluate "Link to this heading")
--------------------------------------------------------------
A common question is when to use [`eval()`](../python/_autosummary/mlx.core.eval.html#mlx.core.eval "mlx.core.eval")
. The trade-off is between letting graphs get too large and not batching enough useful work.
For example:
for \_ in range(100):
a \= a + b
mx.eval(a)
b \= b \* 2
mx.eval(b)
This is a bad idea because there is some fixed overhead with each graph evaluation. On the other hand, there is some slight overhead which grows with the compute graph size, so extremely large graphs (while computationally correct) can be costly.
Luckily, a wide range of compute graph sizes work pretty well with MLX: anything from a few tens of operations to many thousands of operations per evaluation should be okay.
Most numerical computations have an iterative outer loop (e.g. the iteration in stochastic gradient descent). A natural and usually efficient place to use [`eval()`](../python/_autosummary/mlx.core.eval.html#mlx.core.eval "mlx.core.eval")
is at each iteration of this outer loop.
Here is a concrete example:
for batch in dataset:
\# Nothing has been evaluated yet
loss, grad \= value\_and\_grad\_fn(model, batch)
\# Still nothing has been evaluated
optimizer.update(model, grad)
\# Evaluate the loss and the new parameters which will
\# run the full gradient computation and optimizer update
mx.eval(loss, model.parameters())
An important behavior to be aware of is when the graph will be implicitly evaluated. Anytime you `print` an array, convert it to an [`numpy.ndarray`](https://numpy.org/doc/stable/reference/generated/numpy.ndarray.html#numpy.ndarray "(in NumPy v2.2)")
, or otherwise access its memory via [`memoryview`](https://docs.python.org/3/library/stdtypes.html#memoryview "(in Python v3.13)")
, the graph will be evaluated. Saving arrays via [`save()`](../python/_autosummary/mlx.core.save.html#mlx.core.save "mlx.core.save")
(or any other MLX saving functions) will also evaluate the array.
Calling [`array.item()`](../python/_autosummary/mlx.core.array.item.html#mlx.core.array.item "mlx.core.array.item")
on a scalar array will also evaluate it. In the example above, printing the loss (`print(loss)`) or adding the loss scalar to a list (`losses.append(loss.item())`) would cause a graph evaluation. If these lines are before `mx.eval(loss, model.parameters())` then this will be a partial evaluation, computing only the forward pass.
Also, calling [`eval()`](../python/_autosummary/mlx.core.eval.html#mlx.core.eval "mlx.core.eval")
on an array or set of arrays multiple times is perfectly fine. This is effectively a no-op.
Warning
Using scalar arrays for control-flow will cause an evaluation.
Here is an example:
def fun(x):
h, y \= first\_layer(x)
if y \> 0: \# An evaluation is done here!
z \= second\_layer\_a(h)
else:
z \= second\_layer\_b(h)
return z
Using arrays for control flow should be done with care. The above example works and can even be used with gradient transformations. However, this can be very inefficient if evaluations are done too frequently.
Contents
---
# Indexing Arrays — MLX 0.24.0 documentation
[Skip to main content](#main-content)
Back to top
Ctrl+K
[](../index.html)
[](https://github.com/ml-explore/mlx "Source repository")
* [.rst](../_sources/usage/indexing.rst "Download source file")
* .pdf
Indexing Arrays
===============
Contents
--------
Indexing Arrays[#](#indexing-arrays "Link to this heading")
============================================================
For the most part, indexing an MLX [`array`](../python/_autosummary/mlx.core.array.html#mlx.core.array "mlx.core.array")
works the same as indexing a NumPy [`numpy.ndarray`](https://numpy.org/doc/stable/reference/generated/numpy.ndarray.html#numpy.ndarray "(in NumPy v2.2)")
. See the [NumPy documentation](https://numpy.org/doc/stable/user/basics.indexing.html)
for more details on how that works.
For example, you can use regular integers and slices ([`slice`](../python/_autosummary/mlx.core.slice.html#mlx.core.slice "mlx.core.slice")
) to index arrays:
\>>> arr \= mx.arange(10)
>>> arr\[3\]
array(3, dtype\=int32)
>>> arr\[\-2\] \# negative indexing works
array(8, dtype\=int32)
>>> arr\[2:8:2\] \# start, stop, stride
array(\[2, 4, 6\], dtype\=int32)
For multi-dimensional arrays, the `...` or [`Ellipsis`](https://docs.python.org/3/library/constants.html#Ellipsis "(in Python v3.13)")
syntax works as in NumPy:
\>>> arr \= mx.arange(8).reshape(2, 2, 2)
>>> arr\[:, :, 0\]
array(3, dtype\=int32)
array(\[\[0, 2\],\
\[4, 6\]\], dtype\=int32
>>> arr\[..., 0\]
array(\[\[0, 2\],\
\[4, 6\]\], dtype\=int32
You can index with `None` to create a new axis:
\>>> arr \= mx.arange(8)
>>> arr.shape
\[8\]
>>> arr\[None\].shape
\[1, 8\]
You can also use an [`array`](../python/_autosummary/mlx.core.array.html#mlx.core.array "mlx.core.array")
to index another [`array`](../python/_autosummary/mlx.core.array.html#mlx.core.array "mlx.core.array")
:
\>>> arr \= mx.arange(10)
>>> idx \= mx.array(\[5, 7\])
>>> arr\[idx\]
array(\[5, 7\], dtype\=int32)
Mixing and matching integers, [`slice`](../python/_autosummary/mlx.core.slice.html#mlx.core.slice "mlx.core.slice")
, `...`, and [`array`](../python/_autosummary/mlx.core.array.html#mlx.core.array "mlx.core.array")
indices works just as in NumPy.
Other functions which may be useful for indexing arrays are [`take()`](../python/_autosummary/mlx.core.take.html#mlx.core.take "mlx.core.take")
and [`take_along_axis()`](../python/_autosummary/mlx.core.take_along_axis.html#mlx.core.take_along_axis "mlx.core.take_along_axis")
.
Differences from NumPy[#](#differences-from-numpy "Link to this heading")
--------------------------------------------------------------------------
Note
MLX indexing is different from NumPy indexing in two important ways:
* Indexing does not perform bounds checking. Indexing out of bounds is undefined behavior.
* Boolean mask based indexing is not yet supported.
The reason for the lack of bounds checking is that exceptions cannot propagate from the GPU. Performing bounds checking for array indices before launching the kernel would be extremely inefficient.
Indexing with boolean masks is something that MLX may support in the future. In general, MLX has limited support for operations for which output _shapes_ are dependent on input _data_. Other examples of these types of operations which MLX does not yet support include [`numpy.nonzero()`](https://numpy.org/doc/stable/reference/generated/numpy.nonzero.html#numpy.nonzero "(in NumPy v2.2)")
and the single input version of [`numpy.where()`](https://numpy.org/doc/stable/reference/generated/numpy.where.html#numpy.where "(in NumPy v2.2)")
.
In Place Updates[#](#in-place-updates "Link to this heading")
--------------------------------------------------------------
In place updates to indexed arrays are possible in MLX. For example:
\>>> a \= mx.array(\[1, 2, 3\])
>>> a\[2\] \= 0
>>> a
array(\[1, 2, 0\], dtype\=int32)
Just as in NumPy, in place updates will be reflected in all references to the same array:
\>>> a \= mx.array(\[1, 2, 3\])
>>> b \= a
>>> b\[2\] \= 0
>>> b
array(\[1, 2, 0\], dtype\=int32)
>>> a
array(\[1, 2, 0\], dtype\=int32)
Transformations of functions which use in-place updates are allowed and work as expected. For example:
def fun(x, idx):
x\[idx\] \= 2.0
return x.sum()
dfdx \= mx.grad(fun)(mx.array(\[1.0, 2.0, 3.0\]), mx.array(\[1\]))
print(dfdx) \# Prints: array(\[1, 0, 1\], dtype=float32)
In the above `dfdx` will have the correct gradient, namely zeros at `idx` and ones elsewhere.
Contents
---
# Function Transforms — MLX 0.24.0 documentation
[Skip to main content](#main-content)
Back to top
Ctrl+K
[](../index.html)
[](https://github.com/ml-explore/mlx "Source repository")
* [.rst](../_sources/usage/function_transforms.rst "Download source file")
* .pdf
Function Transforms
===================
Contents
--------
Function Transforms[#](#function-transforms "Link to this heading")
====================================================================
MLX uses composable function transformations for automatic differentiation, vectorization, and compute graph optimizations. To see the complete list of function transformations check-out the [API documentation](../python/transforms.html#transforms)
.
The key idea behind composable function transformations is that every transformation returns a function which can be further transformed.
Here is a simple example:
\>>> dfdx \= mx.grad(mx.sin)
>>> dfdx(mx.array(mx.pi))
array(\-1, dtype\=float32)
>>> mx.cos(mx.array(mx.pi))
array(\-1, dtype\=float32)
The output of [`grad()`](../python/_autosummary/mlx.core.grad.html#mlx.core.grad "mlx.core.grad")
on [`sin()`](../python/_autosummary/mlx.core.sin.html#mlx.core.sin "mlx.core.sin")
is simply another function. In this case it is the gradient of the sine function which is exactly the cosine function. To get the second derivative you can do:
\>>> d2fdx2 \= mx.grad(mx.grad(mx.sin))
>>> d2fdx2(mx.array(mx.pi / 2))
array(\-1, dtype\=float32)
>>> mx.sin(mx.array(mx.pi / 2))
array(1, dtype\=float32)
Using [`grad()`](../python/_autosummary/mlx.core.grad.html#mlx.core.grad "mlx.core.grad")
on the output of [`grad()`](../python/_autosummary/mlx.core.grad.html#mlx.core.grad "mlx.core.grad")
is always ok. You keep getting higher order derivatives.
Any of the MLX function transformations can be composed in any order to any depth. See the following sections for more information on [automatic differentiation](#auto-diff)
and [automatic vectorization](#vmap)
. For more information on [`compile()`](../python/_autosummary/mlx.core.compile.html#mlx.core.compile "mlx.core.compile")
see the [compile documentation](compile.html#compile)
.
Automatic Differentiation[#](#automatic-differentiation "Link to this heading")
--------------------------------------------------------------------------------
Automatic differentiation in MLX works on functions rather than on implicit graphs.
Note
If you are coming to MLX from PyTorch, you no longer need functions like `backward`, `zero_grad`, and `detach`, or properties like `requires_grad`.
The most basic example is taking the gradient of a scalar-valued function as we saw above. You can use the [`grad()`](../python/_autosummary/mlx.core.grad.html#mlx.core.grad "mlx.core.grad")
and [`value_and_grad()`](../python/_autosummary/mlx.core.value_and_grad.html#mlx.core.value_and_grad "mlx.core.value_and_grad")
function to compute gradients of more complex functions. By default these functions compute the gradient with respect to the first argument:
def loss\_fn(w, x, y):
return mx.mean(mx.square(w \* x \- y))
w \= mx.array(1.0)
x \= mx.array(\[0.5, \-0.5\])
y \= mx.array(\[1.5, \-1.5\])
\# Computes the gradient of loss\_fn with respect to w:
grad\_fn \= mx.grad(loss\_fn)
dloss\_dw \= grad\_fn(w, x, y)
\# Prints array(-1, dtype=float32)
print(dloss\_dw)
\# To get the gradient with respect to x we can do:
grad\_fn \= mx.grad(loss\_fn, argnums\=1)
dloss\_dx \= grad\_fn(w, x, y)
\# Prints array(\[-1, 1\], dtype=float32)
print(dloss\_dx)
One way to get the loss and gradient is to call `loss_fn` followed by `grad_fn`, but this can result in a lot of redundant work. Instead, you should use [`value_and_grad()`](../python/_autosummary/mlx.core.value_and_grad.html#mlx.core.value_and_grad "mlx.core.value_and_grad")
. Continuing the above example:
\# Computes the gradient of loss\_fn with respect to w:
loss\_and\_grad\_fn \= mx.value\_and\_grad(loss\_fn)
loss, dloss\_dw \= loss\_and\_grad\_fn(w, x, y)
\# Prints array(1, dtype=float32)
print(loss)
\# Prints array(-1, dtype=float32)
print(dloss\_dw)
You can also take the gradient with respect to arbitrarily nested Python containers of arrays (specifically any of [`list`](https://docs.python.org/3/library/stdtypes.html#list "(in Python v3.13)")
, [`tuple`](https://docs.python.org/3/library/stdtypes.html#tuple "(in Python v3.13)")
, or [`dict`](https://docs.python.org/3/library/stdtypes.html#dict "(in Python v3.13)")
).
Suppose we wanted a weight and a bias parameter in the above example. A nice way to do that is the following:
def loss\_fn(params, x, y):
w, b \= params\["weight"\], params\["bias"\]
h \= w \* x + b
return mx.mean(mx.square(h \- y))
params \= {"weight": mx.array(1.0), "bias": mx.array(0.0)}
x \= mx.array(\[0.5, \-0.5\])
y \= mx.array(\[1.5, \-1.5\])
\# Computes the gradient of loss\_fn with respect to both the
\# weight and bias:
grad\_fn \= mx.grad(loss\_fn)
grads \= grad\_fn(params, x, y)
\# Prints
\# {'weight': array(-1, dtype=float32), 'bias': array(0, dtype=float32)}
print(grads)
Notice the tree structure of the parameters is preserved in the gradients.
In some cases you may want to stop gradients from propagating through a part of the function. You can use the [`stop_gradient()`](../python/_autosummary/mlx.core.stop_gradient.html#mlx.core.stop_gradient "mlx.core.stop_gradient")
for that.
Automatic Vectorization[#](#automatic-vectorization "Link to this heading")
----------------------------------------------------------------------------
Use [`vmap()`](../python/_autosummary/mlx.core.vmap.html#mlx.core.vmap "mlx.core.vmap")
to automate vectorizing complex functions. Here we’ll go through a basic and contrived example for the sake of clarity, but [`vmap()`](../python/_autosummary/mlx.core.vmap.html#mlx.core.vmap "mlx.core.vmap")
can be quite powerful for more complex functions which are difficult to optimize by hand.
Warning
Some operations are not yet supported with [`vmap()`](../python/_autosummary/mlx.core.vmap.html#mlx.core.vmap "mlx.core.vmap")
. If you encounter an error like: `ValueError: Primitive's vmap not implemented.` file an [issue](https://github.com/ml-explore/mlx/issues)
and include your function. We will prioritize including it.
A naive way to add the elements from two sets of vectors is with a loop:
xs \= mx.random.uniform(shape\=(4096, 100))
ys \= mx.random.uniform(shape\=(100, 4096))
def naive\_add(xs, ys):
return \[xs\[i\] + ys\[:, i\] for i in range(xs.shape\[0\])\]
Instead you can use [`vmap()`](../python/_autosummary/mlx.core.vmap.html#mlx.core.vmap "mlx.core.vmap")
to automatically vectorize the addition:
\# Vectorize over the second dimension of x and the
\# first dimension of y
vmap\_add \= mx.vmap(lambda x, y: x + y, in\_axes\=(0, 1))
The `in_axes` parameter can be used to specify which dimensions of the corresponding input to vectorize over. Similarly, use `out_axes` to specify where the vectorized axes should be in the outputs.
Let’s time these two different versions:
import timeit
print(timeit.timeit(lambda: mx.eval(naive\_add(xs, ys)), number\=100))
print(timeit.timeit(lambda: mx.eval(vmap\_add(xs, ys)), number\=100))
On an M1 Max the naive version takes in total `5.639` seconds whereas the vectorized version takes only `0.024` seconds, more than 200 times faster.
Of course, this operation is quite contrived. A better approach is to simply do `xs + ys.T`, but for more complex functions [`vmap()`](../python/_autosummary/mlx.core.vmap.html#mlx.core.vmap "mlx.core.vmap")
can be quite handy.
Contents
---
# Using Streams — MLX 0.24.0 documentation
[Skip to main content](#main-content)
Back to top
Ctrl+K
[](../index.html)
[](https://github.com/ml-explore/mlx "Source repository")
* [.rst](../_sources/usage/using_streams.rst "Download source file")
* .pdf
Using Streams
=============
Contents
--------
Using Streams[#](#using-streams "Link to this heading")
========================================================
Specifying the [`Stream`](../python/_autosummary/stream_class.html#mlx.core.Stream "mlx.core.Stream")
[#](#specifying-the-stream "Link to this heading")
---------------------------------------------------------------------------------------------------------------------------------------------------------
All operations (including random number generation) take an optional keyword argument `stream`. The `stream` kwarg specifies which [`Stream`](../python/_autosummary/stream_class.html#mlx.core.Stream "mlx.core.Stream")
the operation should run on. If the stream is unspecified then the operation is run on the default stream of the default device: `mx.default_stream(mx.default_device())`. The `stream` kwarg can also be a [`Device`](../python/_autosummary/mlx.core.Device.html#mlx.core.Device "mlx.core.Device")
(e.g. `stream=my_device`) in which case the operation is run on the default stream of the provided device `mx.default_stream(my_device)`.
Contents
---
# Compilation — MLX 0.24.0 documentation
[Skip to main content](#main-content)
Back to top
Ctrl+K
[](../index.html)
[](https://github.com/ml-explore/mlx "Source repository")
* [.rst](../_sources/usage/compile.rst "Download source file")
* .pdf
Compilation
===========
Contents
--------
Compilation[#](#compilation "Link to this heading")
====================================================
MLX has a [`compile()`](../python/_autosummary/mlx.core.compile.html#mlx.core.compile "mlx.core.compile")
function transformation which compiles computation graphs. Function compilation results in smaller graphs by merging common work and fusing certain operations. In many cases this can lead to big improvements in run-time and memory use.
Getting started with [`compile()`](../python/_autosummary/mlx.core.compile.html#mlx.core.compile "mlx.core.compile")
is simple, but there are some edge cases that are good to be aware of for more complex graphs and advanced usage.
Basics of Compile[#](#basics-of-compile "Link to this heading")
----------------------------------------------------------------
Let’s start with a simple example:
def fun(x, y):
return mx.exp(\-x) + y
x \= mx.array(1.0)
y \= mx.array(2.0)
\# Regular call, no compilation
\# Prints: array(2.36788, dtype=float32)
print(fun(x, y))
\# Compile the function
compiled\_fun \= mx.compile(fun)
\# Prints: array(2.36788, dtype=float32)
print(compiled\_fun(x, y))
The output of both the regular function and the compiled function is the same up to numerical precision.
The first time you call a compiled function, MLX will build the compute graph, optimize it, and generate and compile code. This can be relatively slow. However, MLX will cache compiled functions, so calling a compiled function multiple times will not initiate a new compilation. This means you should typically compile functions that you plan to use more than once.
def fun(x, y):
return mx.exp(\-x) + y
x \= mx.array(1.0)
y \= mx.array(2.0)
compiled\_fun \= mx.compile(fun)
\# Compiled here
compiled\_fun(x, y)
\# Not compiled again
compiled\_fun(x, y)
\# Not compiled again
mx.compile(fun)(x, y)
There are some important cases to be aware of that can cause a function to be recompiled:
* Changing the shape or number of dimensions
* Changing the type of any of the inputs
* Changing the number of inputs to the function
In certain cases only some of the compilation stack will be rerun (for example when changing the shapes) and in other cases the full compilation stack will be rerun (for example when changing the types). In general you should avoid compiling functions too frequently.
Another idiom to watch out for is compiling functions which get created and destroyed frequently. This can happen, for example, when compiling an anonymous function in a loop:
a \= mx.array(1.0)
\# Don't do this, compiles lambda at each iteration
for \_ in range(5):
mx.compile(lambda x: mx.exp(mx.abs(x)))(a)
Example Speedup[#](#example-speedup "Link to this heading")
------------------------------------------------------------
The [`mlx.nn.gelu()`](../python/nn/_autosummary_functions/mlx.nn.gelu.html#mlx.nn.gelu "mlx.nn.gelu")
is a nonlinear activation function commonly used with Transformer-based models. The implementation involves several unary and binary element-wise operations:
def gelu(x):
return x \* (1 + mx.erf(x / math.sqrt(2))) / 2
If you use this function with small arrays, it will be overhead bound. If you use it with large arrays it will be memory bandwidth bound. However, all of the operations in the `gelu` are fusible into a single kernel with [`compile()`](../python/_autosummary/mlx.core.compile.html#mlx.core.compile "mlx.core.compile")
. This can speedup both cases considerably.
Let’s compare the runtime of the regular function versus the compiled function. We’ll use the following timing helper which does a warm up and handles synchronization:
import time
def timeit(fun, x):
\# warm up
for \_ in range(10):
mx.eval(fun(x))
tic \= time.perf\_counter()
for \_ in range(100):
mx.eval(fun(x))
toc \= time.perf\_counter()
tpi \= 1e3 \* (toc \- tic) / 100
print(f"Time per iteration {tpi:.3f} (ms)")
Now make an array, and benchmark both functions:
x \= mx.random.uniform(shape\=(32, 1000, 4096))
timeit(nn.gelu, x)
timeit(mx.compile(nn.gelu), x)
On an M1 Max the times are 15.5 and 3.1 milliseconds. The compiled `gelu` is five times faster.
Debugging[#](#debugging "Link to this heading")
------------------------------------------------
When a compiled function is first called, it is traced with placeholder inputs. This means you can’t evaluate arrays (for example to print their contents) inside compiled functions.
@mx.compile
def fun(x):
z \= \-x
print(z) \# Crash
return mx.exp(z)
fun(mx.array(5.0))
For debugging, inspecting arrays can be helpful. One way to do that is to globally disable compilation using the [`disable_compile()`](../python/_autosummary/mlx.core.disable_compile.html#mlx.core.disable_compile "mlx.core.disable_compile")
function or `MLX_DISABLE_COMPILE` flag. For example the following is okay even though `fun` is compiled:
@mx.compile
def fun(x):
z \= \-x
print(z) \# Okay
return mx.exp(z)
mx.disable\_compile()
fun(mx.array(5.0))
Pure Functions[#](#pure-functions "Link to this heading")
----------------------------------------------------------
Compiled functions are intended to be _pure_; that is they should not have side effects. For example:
state \= \[\]
@mx.compile
def fun(x, y):
z \= x + y
state.append(z)
return mx.exp(z)
fun(mx.array(1.0), mx.array(2.0))
\# Crash!
print(state)
After the first call of `fun`, the `state` list will hold a placeholder array. The placeholder does not have any data; it is only used to build the computation graph. Printing such an array results in a crash.
You have two options to deal with this. The first option is to simply return `state` as an output:
state \= \[\]
@mx.compile
def fun(x, y):
z \= x + y
state.append(z)
return mx.exp(z), state
\_, state \= fun(mx.array(1.0), mx.array(2.0))
\# Prints \[array(3, dtype=float32)\]
print(state)
In some cases returning updated state can be pretty inconvenient. Hence, [`compile()`](../python/_autosummary/mlx.core.compile.html#mlx.core.compile "mlx.core.compile")
has a parameter to capture implicit outputs:
from functools import partial
state \= \[\]
\# Tell compile to capture state as an output
@partial(mx.compile, outputs\=state)
def fun(x, y):
z \= x + y
state.append(z)
return mx.exp(z), state
fun(mx.array(1.0), mx.array(2.0))
\# Prints \[array(3, dtype=float32)\]
print(state)
This is particularly useful for compiling a function which includes an update to a container of arrays, as is commonly done when training the parameters of a [`mlx.nn.Module`](../python/nn/module.html#mlx.nn.Module "mlx.nn.Module")
.
Compiled functions will also treat any inputs not in the parameter list as constants. For example:
state \= \[mx.array(1.0)\]
@mx.compile
def fun(x):
return x + state\[0\]
\# Prints array(2, dtype=float32)
print(fun(mx.array(1.0)))
\# Update state
state\[0\] \= mx.array(5.0)
\# Still prints array(2, dtype=float32)
print(fun(mx.array(1.0)))
In order to have the change of state reflected in the outputs of `fun` you again have two options. The first option is to simply pass `state` as input to the function. In some cases this can be pretty inconvenient. Hence, [`compile()`](../python/_autosummary/mlx.core.compile.html#mlx.core.compile "mlx.core.compile")
also has a parameter to capture implicit inputs:
from functools import partial
state \= \[mx.array(1.0)\]
\# Tell compile to capture state as an input
@partial(mx.compile, inputs\=state)
def fun(x):
return x + state\[0\]
\# Prints array(2, dtype=float32)
print(fun(mx.array(1.0)))
\# Update state
state\[0\] \= mx.array(5.0)
\# Prints array(6, dtype=float32)
print(fun(mx.array(1.0)))
Compiling Training Graphs[#](#compiling-training-graphs "Link to this heading")
--------------------------------------------------------------------------------
This section will step through how to use [`compile()`](../python/_autosummary/mlx.core.compile.html#mlx.core.compile "mlx.core.compile")
with a simple example of a common setup: training a model with [`mlx.nn.Module`](../python/nn/module.html#mlx.nn.Module "mlx.nn.Module")
using an [`mlx.optimizers.Optimizer`](../python/optimizers/optimizer.html#mlx.optimizers.Optimizer "mlx.optimizers.Optimizer")
with state. We will show how to compile the full forward, backward, and update with [`compile()`](../python/_autosummary/mlx.core.compile.html#mlx.core.compile "mlx.core.compile")
.
To start, here is the simple example without any compilation:
import mlx.core as mx
import mlx.nn as nn
import mlx.optimizers as optim
\# 4 examples with 10 features each
x \= mx.random.uniform(shape\=(4, 10))
\# 0, 1 targets
y \= mx.array(\[0, 1, 0, 1\])
\# Simple linear model
model \= nn.Linear(10, 1)
\# SGD with momentum
optimizer \= optim.SGD(learning\_rate\=0.1, momentum\=0.8)
def loss\_fn(model, x, y):
logits \= model(x).squeeze()
return nn.losses.binary\_cross\_entropy(logits, y)
loss\_and\_grad\_fn \= nn.value\_and\_grad(model, loss\_fn)
\# Perform 10 steps of gradient descent
for it in range(10):
loss, grads \= loss\_and\_grad\_fn(model, x, y)
optimizer.update(model, grads)
mx.eval(model.parameters(), optimizer.state)
To compile the update we can put it all in a function and compile it with the appropriate input and output captures. Here’s the same example but compiled:
import mlx.core as mx
import mlx.nn as nn
import mlx.optimizers as optim
from functools import partial
\# 4 examples with 10 features each
x \= mx.random.uniform(shape\=(4, 10))
\# 0, 1 targets
y \= mx.array(\[0, 1, 0, 1\])
\# Simple linear model
model \= nn.Linear(10, 1)
\# SGD with momentum
optimizer \= optim.SGD(learning\_rate\=0.1, momentum\=0.8)
def loss\_fn(model, x, y):
logits \= model(x).squeeze()
return nn.losses.binary\_cross\_entropy(logits, y)
\# The state that will be captured as input and output
state \= \[model.state, optimizer.state\]
@partial(mx.compile, inputs\=state, outputs\=state)
def step(x, y):
loss\_and\_grad\_fn \= nn.value\_and\_grad(model, loss\_fn)
loss, grads \= loss\_and\_grad\_fn(model, x, y)
optimizer.update(model, grads)
return loss
\# Perform 10 steps of gradient descent
for it in range(10):
loss \= step(x, y)
\# Evaluate the model and optimizer state
mx.eval(state)
print(loss)
Note
If you are using a module which performs random sampling such as [`mlx.nn.Dropout()`](../python/nn/_autosummary/mlx.nn.Dropout.html#mlx.nn.Dropout "mlx.nn.Dropout")
, make sure you also include `mx.random.state` in the `state` captured by [`compile()`](../python/_autosummary/mlx.core.compile.html#mlx.core.compile "mlx.core.compile")
, i.e. `state = [model.state, optimizer.state, mx.random.state]`.
Note
For more examples of compiling full training graphs checkout the [MLX Examples](https://github.com/ml-explore/mlx-examples)
GitHub repo.
Transformations with Compile[#](#transformations-with-compile "Link to this heading")
--------------------------------------------------------------------------------------
In MLX function transformations are composable. You can apply any function transformation to the output of any other function transformation. For more on this, see the documentation on [function transforms](function_transforms.html#function-transforms)
.
Compiling transformed functions works just as expected:
grad\_fn \= mx.grad(mx.exp)
compiled\_grad\_fn \= mx.compile(grad\_fn)
\# Prints: array(2.71828, dtype=float32)
print(grad\_fn(mx.array(1.0)))
\# Also prints: array(2.71828, dtype=float32)
print(compiled\_grad\_fn(mx.array(1.0)))
Note
In order to compile as much as possible, a transformation of a compiled function will not by default be compiled. To compile the transformed function simply pass it through [`compile()`](../python/_autosummary/mlx.core.compile.html#mlx.core.compile "mlx.core.compile")
.
You can also compile functions which themselves call compiled functions. A good practice is to compile the outer most function to give [`compile()`](../python/_autosummary/mlx.core.compile.html#mlx.core.compile "mlx.core.compile")
the most opportunity to optimize the computation graph:
@mx.compile
def inner(x):
return mx.exp(\-mx.abs(x))
def outer(x):
inner(inner(x))
\# Compiling the outer function is good to do as it will likely
\# be faster even though the inner functions are compiled
fun \= mx.compile(outer)
Shapeless Compilation[#](#shapeless-compilation "Link to this heading")
------------------------------------------------------------------------
When the shape of an input to a compiled function changes, the function is recompiled. You can compile a function once and run it on inputs with variable shapes by specifying `shapeless=True` to [`compile()`](../python/_autosummary/mlx.core.compile.html#mlx.core.compile "mlx.core.compile")
. In this case changes to the shapes of the inputs do not cause the function to be recompiled.
def fun(x, y):
return mx.abs(x + y)
compiled\_fun \= mx.compile(fun, shapeless\=True)
x \= mx.array(1.0)
y \= mx.array(\-2.0)
\# Firt call compiles the function
print(compiled\_fun(x, y))
\# Second call with different shapes
\# does not recompile the function
x \= mx.array(\[1.0, \-6.0\])
y \= mx.array(\[\-2.0, 3.0\])
print(compiled\_fun(x, y))
Use shapeless compilations carefully. Since compilation is not triggered when shapes change, any graphs which are conditional on the input shapes will not work as expected. Shape-dependent computations are common and sometimes subtle to detect. For example:
def fun(x):
return x.reshape(x.shape\[0\] \* x.shape\[1\], \-1)
compiled\_fun \= mx.compile(fun, shapeless\=True)
x \= mx.random.uniform(shape\=(2, 3, 4))
out \= compiled\_fun(x)
x \= mx.random.uniform(shape\=(5, 5, 3))
\# Error, can't reshape (5, 5, 3) to (6, -1)
out \= compiled\_fun(x)
The second call to the `compiled_fun` fails because of the call to [`reshape()`](../python/_autosummary/mlx.core.reshape.html#mlx.core.reshape "mlx.core.reshape")
which uses the static shape of `x` in the first call. We can fix this by using [`flatten()`](../python/_autosummary/mlx.core.flatten.html#mlx.core.flatten "mlx.core.flatten")
to avoid hardcoding the shape of `x`:
def fun(x):
return x.flatten(0, 1)
compiled\_fun \= mx.compile(fun, shapeless\=True)
x \= mx.random.uniform(shape\=(2, 3, 4))
out \= compiled\_fun(x)
x \= mx.random.uniform(shape\=(5, 5, 3))
\# Ok
out \= compiled\_fun(x)
Contents
---
# Exporting Functions — MLX 0.24.0 documentation
[Skip to main content](#main-content)
Back to top
Ctrl+K
[](../index.html)
[](https://github.com/ml-explore/mlx "Source repository")
* [.rst](../_sources/usage/export.rst "Download source file")
* .pdf
Exporting Functions
===================
Contents
--------
Exporting Functions[#](#exporting-functions "Link to this heading")
====================================================================
MLX has an API to export and import functions to and from a file. This lets you run computations written in one MLX front-end (e.g. Python) in another MLX front-end (e.g. C++).
This guide walks through the basics of the MLX export API with some examples. To see the full list of functions check-out the [API documentation](../python/export.html#export)
.
Basics of Exporting[#](#basics-of-exporting "Link to this heading")
--------------------------------------------------------------------
Let’s start with a simple example:
def fun(x, y):
return x + y
x \= mx.array(1.0)
y \= mx.array(1.0)
mx.export\_function("add.mlxfn", fun, x, y)
To export a function, provide sample input arrays that the function can be called with. The data doesn’t matter, but the shapes and types of the arrays do. In the above example we exported `fun` with two `float32` scalar arrays. We can then import the function and run it:
add\_fun \= mx.import\_function("add.mlxfn")
out, \= add\_fun(mx.array(1.0), mx.array(2.0))
\# Prints: array(3, dtype=float32)
print(out)
out, \= add\_fun(mx.array(1.0), mx.array(3.0))
\# Prints: array(4, dtype=float32)
print(out)
\# Raises an exception
add\_fun(mx.array(1), mx.array(3.0))
\# Raises an exception
add\_fun(mx.array(\[1.0, 2.0\]), mx.array(3.0))
Notice the third and fourth calls to `add_fun` raise exceptions because the shapes and types of the inputs are different than the shapes and types of the example inputs we exported the function with.
Also notice that even though the original `fun` returns a single output array, the imported function always returns a tuple of one or more arrays.
The inputs to [`export_function()`](../python/_autosummary/mlx.core.export_function.html#mlx.core.export_function "mlx.core.export_function")
and to an imported function can be specified as variable positional arguments or as a tuple of arrays:
def fun(x, y):
return x + y
x \= mx.array(1.0)
y \= mx.array(1.0)
\# Both arguments to fun are positional
mx.export\_function("add.mlxfn", fun, x, y)
\# Same as above
mx.export\_function("add.mlxfn", fun, (x, y))
imported\_fun \= mx.import\_function("add.mlxfn")
\# Ok
out, \= imported\_fun(x, y)
\# Also ok
out, \= imported\_fun((x, y))
You can pass example inputs to functions as positional or keyword arguments. If you use keyword arguments to export the function, then you have to use the same keyword arguments when calling the imported function.
def fun(x, y):
return x + y
\# One argument to fun is positional, the other is a kwarg
mx.export\_function("add.mlxfn", fun, x, y\=y)
imported\_fun \= mx.import\_function("add.mlxfn")
\# Ok
out, \= imported\_fun(x, y\=y)
\# Also ok
out, \= imported\_fun((x,), {"y": y})
\# Raises since the keyword argument is missing
out, \= imported\_fun(x, y)
\# Raises since the keyword argument has the wrong key
out, \= imported\_fun(x, z\=y)
Exporting Modules[#](#exporting-modules "Link to this heading")
----------------------------------------------------------------
An [`mlx.nn.Module`](../python/nn/module.html#mlx.nn.Module "mlx.nn.Module")
can be exported with or without the parameters included in the exported function. Here’s an example:
model \= nn.Linear(4, 4)
mx.eval(model.parameters())
def call(x):
return model(x)
mx.export\_function("model.mlxfn", call, mx.zeros(4))
In the above example, the [`mlx.nn.Linear`](../python/nn/_autosummary/mlx.nn.Linear.html#mlx.nn.Linear "mlx.nn.Linear")
module is exported. Its parameters are also saved to the `model.mlxfn` file.
Note
For enclosed arrays inside an exported function, be extra careful to ensure they are evaluated. The computation graph that gets exported will include the computation that produces enclosed inputs.
If the above example was missing `mx.eval(model.parameters()`, the exported function would include the random initialization of the [`mlx.nn.Module`](../python/nn/module.html#mlx.nn.Module "mlx.nn.Module")
parameters.
If you only want to export the `Module.__call__` function without the parameters, pass them as inputs to the `call` wrapper:
model \= nn.Linear(4, 4)
mx.eval(model.parameters())
def call(x, \*\*params):
\# Set the model's parameters to the input parameters
model.update(tree\_unflatten(list(params.items())))
return model(x)
params \= dict(tree\_flatten(model.parameters()))
mx.export\_function("model.mlxfn", call, (mx.zeros(4),), params)
Shapeless Exports[#](#shapeless-exports "Link to this heading")
----------------------------------------------------------------
Just like [`compile()`](../python/_autosummary/mlx.core.compile.html#mlx.core.compile "mlx.core.compile")
, functions can also be exported for dynamically shaped inputs. Pass `shapeless=True` to [`export_function()`](../python/_autosummary/mlx.core.export_function.html#mlx.core.export_function "mlx.core.export_function")
or [`exporter()`](../python/_autosummary/mlx.core.exporter.html#mlx.core.exporter "mlx.core.exporter")
to export a function which can be used for inputs with variable shapes:
mx.export\_function("fun.mlxfn", mx.abs, mx.array(0.0), shapeless\=True)
imported\_abs \= mx.import\_function("fun.mlxfn")
\# Ok
out, \= imported\_abs(mx.array(\-1.0))
\# Also ok
out, \= imported\_abs(mx.array(\[\-1.0, \-2.0\]))
With `shapeless=False` (which is the default), the second call to `imported_abs` would raise an exception with a shape mismatch.
Shapeless exporting works the same as shapeless compilation and should be used carefully. See the [documentation on shapeless compilation](compile.html#shapeless-compile)
for more information.
Exporting Multiple Traces[#](#exporting-multiple-traces "Link to this heading")
--------------------------------------------------------------------------------
In some cases, functions build different computation graphs for different input arguments. A simple way to manage this is to export to a new file with each set of inputs. This is a fine option in many cases. But it can be suboptimal if the exported functions have a large amount of duplicate constant data (for example the parameters of a [`mlx.nn.Module`](../python/nn/module.html#mlx.nn.Module "mlx.nn.Module")
).
The export API in MLX lets you export multiple traces of the same function to a single file by creating an exporting context manager with [`exporter()`](../python/_autosummary/mlx.core.exporter.html#mlx.core.exporter "mlx.core.exporter")
:
def fun(x, y\=None):
constant \= mx.array(3.0)
if y is not None:
x += y
return x + constant
with mx.exporter("fun.mlxfn", fun) as exporter:
exporter(mx.array(1.0))
exporter(mx.array(1.0), y\=mx.array(0.0))
imported\_function \= mx.import\_function("fun.mlxfn")
\# Call the function with y=None
out, \= imported\_function(mx.array(1.0))
print(out)
\# Call the function with y specified
out, \= imported\_function(mx.array(1.0), y\=mx.array(1.0))
print(out)
In the above example the function constant data, (i.e. `constant`), is only saved once.
Transformations with Imported Functions[#](#transformations-with-imported-functions "Link to this heading")
------------------------------------------------------------------------------------------------------------
Function transformations like [`grad()`](../python/_autosummary/mlx.core.grad.html#mlx.core.grad "mlx.core.grad")
, [`vmap()`](../python/_autosummary/mlx.core.vmap.html#mlx.core.vmap "mlx.core.vmap")
, and [`compile()`](../python/_autosummary/mlx.core.compile.html#mlx.core.compile "mlx.core.compile")
work on imported functions just like regular Python functions:
def fun(x):
return mx.sin(x)
x \= mx.array(0.0)
mx.export\_function("sine.mlxfn", fun, x)
imported\_fun \= mx.import\_function("sine.mlxfn")
\# Take the derivative of the imported function
dfdx \= mx.grad(lambda x: imported\_fun(x)\[0\])
\# Prints: array(1, dtype=float32)
print(dfdx(x))
\# Compile the imported function
mx.compile(imported\_fun)
\# Prints: array(0, dtype=float32)
print(compiled\_fun(x)\[0\])
Importing Functions in C++[#](#importing-functions-in-c "Link to this heading")
--------------------------------------------------------------------------------
Importing and running functions in C++ is basically the same as importing and running them in Python. First, follow the [instructions](../dev/mlx_in_cpp.html#mlx-in-cpp)
to setup a simple C++ project that uses MLX as a library.
Next, export a simple function from Python:
def fun(x, y):
return mx.exp(x + y)
x \= mx.array(1.0)
y \= mx.array(1.0)
mx.export\_function("fun.mlxfn", fun, x, y)
Import and run the function in C++ with only a few lines of code:
auto fun \= mx::import\_function("fun.mlxfn");
auto inputs \= {mx::array(1.0), mx::array(1.0)};
auto outputs \= fun(inputs);
// Prints: array(2, dtype=float32)
std::cout << outputs\[0\] << std::endl;
Imported functions can be transformed in C++ just like in Python. Use `std::vector` for positional arguments and `std::map` for keyword arguments when calling imported functions in C++.
More Examples[#](#more-examples "Link to this heading")
--------------------------------------------------------
Here are a few more complete examples exporting more complex functions from Python and importing and running them in C++:
* [Inference and training a multi-layer perceptron](https://github.com/ml-explore/mlx/tree/main/examples/export)
Contents
---
# Conversion to NumPy and Other Frameworks — MLX 0.24.0 documentation
[Skip to main content](#main-content)
Back to top
Ctrl+K
[](../index.html)
[](https://github.com/ml-explore/mlx "Source repository")
* [.rst](../_sources/usage/numpy.rst "Download source file")
* .pdf
Conversion to NumPy and Other Frameworks
========================================
Contents
--------
Conversion to NumPy and Other Frameworks[#](#conversion-to-numpy-and-other-frameworks "Link to this heading")
==============================================================================================================
MLX array supports conversion between other frameworks with either:
* The [Python Buffer Protocol](https://docs.python.org/3/c-api/buffer.html)
.
* [DLPack](https://dmlc.github.io/dlpack/latest/)
.
Let’s convert an array to NumPy and back.
import mlx.core as mx
import numpy as np
a \= mx.arange(3)
b \= np.array(a) \# copy of a
c \= mx.array(b) \# copy of b
Note
Since NumPy does not support `bfloat16` arrays, you will need to convert to `float16` or `float32` first: `np.array(a.astype(mx.float32))`. Otherwise, you will receive an error like: `Item size 2 for PEP 3118 buffer format string does not match the dtype V item size 0.`
By default, NumPy copies data to a new array. This can be prevented by creating an array view:
a \= mx.arange(3)
a\_view \= np.array(a, copy\=False)
print(a\_view.flags.owndata) \# False
a\_view\[0\] \= 1
print(a\[0\].item()) \# 1
Note
NumPy arrays with type `float64` will be default converted to MLX arrays with type `float32`.
A NumPy array view is a normal NumPy array, except that it does not own its memory. This means writing to the view is reflected in the original array.
While this is quite powerful to prevent copying arrays, it should be noted that external changes to the memory of arrays cannot be reflected in gradients.
Let’s demonstrate this in an example:
def f(x):
x\_view \= np.array(x, copy\=False)
x\_view\[:\] \*= x\_view \# modify memory without telling mx
return x.sum()
x \= mx.array(\[3.0\])
y, df \= mx.value\_and\_grad(f)(x)
print("f(x) = x² =", y.item()) \# 9.0
print("f'(x) = 2x !=", df.item()) \# 1.0
The function `f` indirectly modifies the array `x` through a memory view. However, this modification is not reflected in the gradient, as seen in the last line outputting `1.0`, representing the gradient of the sum operation alone. The squaring of `x` occurs externally to MLX, meaning that no gradient is incorporated. It’s important to note that a similar issue arises during array conversion and copying. For instance, a function defined as `mx.array(np.array(x)**2).sum()` would also result in an incorrect gradient, even though no in-place operations on MLX memory are executed.
PyTorch[#](#pytorch "Link to this heading")
--------------------------------------------
Warning
PyTorch Support for [`memoryview`](https://docs.python.org/3/library/stdtypes.html#memoryview "(in Python v3.13)")
is experimental and can break for multi-dimensional arrays. Casting to NumPy first is advised for now.
PyTorch supports the buffer protocol, but it requires an explicit [`memoryview`](https://docs.python.org/3/library/stdtypes.html#memoryview "(in Python v3.13)")
.
import mlx.core as mx
import torch
a \= mx.arange(3)
b \= torch.tensor(memoryview(a))
c \= mx.array(b.numpy())
Conversion from PyTorch tensors back to arrays must be done via intermediate NumPy arrays with `numpy()`.
JAX[#](#jax "Link to this heading")
------------------------------------
JAX fully supports the buffer protocol.
import mlx.core as mx
import jax.numpy as jnp
a \= mx.arange(3)
b \= jnp.array(a)
c \= mx.array(b)
TensorFlow[#](#tensorflow "Link to this heading")
--------------------------------------------------
TensorFlow supports the buffer protocol, but it requires an explicit [`memoryview`](https://docs.python.org/3/library/stdtypes.html#memoryview "(in Python v3.13)")
.
import mlx.core as mx
import tensorflow as tf
a \= mx.arange(3)
b \= tf.constant(memoryview(a))
c \= mx.array(b)
Contents
---
# Linear Regression — MLX 0.24.0 documentation
[Skip to main content](#main-content)
Back to top
Ctrl+K
[](../index.html)
[](https://github.com/ml-explore/mlx "Source repository")
* [.rst](../_sources/examples/linear_regression.rst "Download source file")
* .pdf
Linear Regression
=================
Linear Regression[#](#linear-regression "Link to this heading")
================================================================
Let’s implement a basic linear regression model as a starting point to learn MLX. First import the core package and setup some problem metadata:
import mlx.core as mx
num\_features \= 100
num\_examples \= 1\_000
num\_iters \= 10\_000 \# iterations of SGD
lr \= 0.01 \# learning rate for SGD
We’ll generate a synthetic dataset by:
1. Sampling the design matrix `X`.
2. Sampling a ground truth parameter vector `w_star`.
3. Compute the dependent values `y` by adding Gaussian noise to `X @ w_star`.
\# True parameters
w\_star \= mx.random.normal((num\_features,))
\# Input examples (design matrix)
X \= mx.random.normal((num\_examples, num\_features))
\# Noisy labels
eps \= 1e-2 \* mx.random.normal((num\_examples,))
y \= X @ w\_star + eps
We will use SGD to find the optimal weights. To start, define the squared loss and get the gradient function of the loss with respect to the parameters.
def loss\_fn(w):
return 0.5 \* mx.mean(mx.square(X @ w \- y))
grad\_fn \= mx.grad(loss\_fn)
Start the optimization by initializing the parameters `w` randomly. Then repeatedly update the parameters for `num_iters` iterations.
w \= 1e-2 \* mx.random.normal((num\_features,))
for \_ in range(num\_iters):
grad \= grad\_fn(w)
w \= w \- lr \* grad
mx.eval(w)
Finally, compute the loss of the learned parameters and verify that they are close to the ground truth parameters.
loss \= loss\_fn(w)
error\_norm \= mx.sum(mx.square(w \- w\_star)).item() \*\* 0.5
print(
f"Loss {loss.item():.5f}, |w-w\*| = {error\_norm:.5f}, "
)
\# Should print something close to: Loss 0.00005, |w-w\*| = 0.00364
Complete [linear regression](https://github.com/ml-explore/mlx/tree/main/examples/python/linear_regression.py)
and [logistic regression](https://github.com/ml-explore/mlx/tree/main/examples/python/logistic_regression.py)
examples are available in the MLX GitHub repo.
---
# Multi-Layer Perceptron — MLX 0.24.0 documentation
[Skip to main content](#main-content)
Back to top
Ctrl+K
[](../index.html)
[](https://github.com/ml-explore/mlx "Source repository")
* [.rst](../_sources/examples/mlp.rst "Download source file")
* .pdf
Multi-Layer Perceptron
======================
Multi-Layer Perceptron[#](#multi-layer-perceptron "Link to this heading")
==========================================================================
In this example we’ll learn to use `mlx.nn` by implementing a simple multi-layer perceptron to classify MNIST.
As a first step import the MLX packages we need:
import mlx.core as mx
import mlx.nn as nn
import mlx.optimizers as optim
import numpy as np
The model is defined as the `MLP` class which inherits from [`mlx.nn.Module`](../python/nn/module.html#mlx.nn.Module "mlx.nn.Module")
. We follow the standard idiom to make a new module:
1. Define an `__init__` where the parameters and/or submodules are setup. See the [Module class docs](../python/nn.html#module-class)
for more information on how [`mlx.nn.Module`](../python/nn/module.html#mlx.nn.Module "mlx.nn.Module")
registers parameters.
2. Define a `__call__` where the computation is implemented.
class MLP(nn.Module):
def \_\_init\_\_(
self, num\_layers: int, input\_dim: int, hidden\_dim: int, output\_dim: int
):
super().\_\_init\_\_()
layer\_sizes \= \[input\_dim\] + \[hidden\_dim\] \* num\_layers + \[output\_dim\]
self.layers \= \[\
nn.Linear(idim, odim)\
for idim, odim in zip(layer\_sizes\[:\-1\], layer\_sizes\[1:\])\
\]
def \_\_call\_\_(self, x):
for l in self.layers\[:\-1\]:
x \= mx.maximum(l(x), 0.0)
return self.layers\[\-1\](x)
We define the loss function which takes the mean of the per-example cross entropy loss. The `mlx.nn.losses` sub-package has implementations of some commonly used loss functions.
def loss\_fn(model, X, y):
return mx.mean(nn.losses.cross\_entropy(model(X), y))
We also need a function to compute the accuracy of the model on the validation set:
def eval\_fn(model, X, y):
return mx.mean(mx.argmax(model(X), axis\=1) \== y)
Next, setup the problem parameters and load the data. To load the data, you need our [mnist data loader](https://github.com/ml-explore/mlx-examples/blob/main/mnist/mnist.py)
, which we will import as `mnist`.
num\_layers \= 2
hidden\_dim \= 32
num\_classes \= 10
batch\_size \= 256
num\_epochs \= 10
learning\_rate \= 1e-1
\# Load the data
import mnist
train\_images, train\_labels, test\_images, test\_labels \= map(
mx.array, mnist.mnist()
)
Since we’re using SGD, we need an iterator which shuffles and constructs minibatches of examples in the training set:
def batch\_iterate(batch\_size, X, y):
perm \= mx.array(np.random.permutation(y.size))
for s in range(0, y.size, batch\_size):
ids \= perm\[s : s + batch\_size\]
yield X\[ids\], y\[ids\]
Finally, we put it all together by instantiating the model, the [`mlx.optimizers.SGD`](../python/optimizers/_autosummary/mlx.optimizers.SGD.html#mlx.optimizers.SGD "mlx.optimizers.SGD")
optimizer, and running the training loop:
\# Load the model
model \= MLP(num\_layers, train\_images.shape\[\-1\], hidden\_dim, num\_classes)
mx.eval(model.parameters())
\# Get a function which gives the loss and gradient of the
\# loss with respect to the model's trainable parameters
loss\_and\_grad\_fn \= nn.value\_and\_grad(model, loss\_fn)
\# Instantiate the optimizer
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 optimizer state and model parameters
\# in a single call
optimizer.update(model, grads)
\# Force a graph evaluation
mx.eval(model.parameters(), optimizer.state)
accuracy \= eval\_fn(model, test\_images, test\_labels)
print(f"Epoch {e}: Test accuracy {accuracy.item():.3f}")
Note
The [`mlx.nn.value_and_grad()`](../python/_autosummary/mlx.nn.value_and_grad.html#mlx.nn.value_and_grad "mlx.nn.value_and_grad")
function is a convenience function to get the gradient of a loss with respect to the trainable parameters of a model. This should not be confused with [`mlx.core.value_and_grad()`](../python/_autosummary/mlx.core.value_and_grad.html#mlx.core.value_and_grad "mlx.core.value_and_grad")
.
The model should train to a decent accuracy (about 95%) after just a few passes over the training set. The [full example](https://github.com/ml-explore/mlx-examples/tree/main/mnist)
is available in the MLX GitHub repo.
---
# LLM inference — MLX 0.24.0 documentation
[Skip to main content](#main-content)
Back to top
Ctrl+K
[](../index.html)
[](https://github.com/ml-explore/mlx "Source repository")
* [.rst](../_sources/examples/llama-inference.rst "Download source file")
* .pdf
LLM inference
=============
Contents
--------
LLM inference[#](#llm-inference "Link to this heading")
========================================================
MLX enables efficient inference of large-ish transformers on Apple silicon without compromising on ease of use. In this example we will create an inference script for the Llama family of transformer models in which the model is defined in less than 200 lines of python.
Implementing the model[#](#implementing-the-model "Link to this heading")
--------------------------------------------------------------------------
We will use the neural network building blocks defined in the `mlx.nn` module to concisely define the model architecture.
### Attention layer[#](#attention-layer "Link to this heading")
We will start with the Llama attention layer which notably uses the RoPE positional encoding. [\[1\]](#id4)
In addition, our attention layer will optionally use a key/value cache that will be concatenated with the provided keys and values to support efficient inference.
Our implementation uses [`mlx.nn.Linear`](../python/nn/_autosummary/mlx.nn.Linear.html#mlx.nn.Linear "mlx.nn.Linear")
for all the projections and [`mlx.nn.RoPE`](../python/nn/_autosummary/mlx.nn.RoPE.html#mlx.nn.RoPE "mlx.nn.RoPE")
for the positional encoding.
import mlx.core as mx
import mlx.nn as nn
class LlamaAttention(nn.Module):
def \_\_init\_\_(self, dims: int, num\_heads: int):
super().\_\_init\_\_()
self.num\_heads \= num\_heads
self.rope \= nn.RoPE(dims // num\_heads, traditional\=True)
self.query\_proj \= nn.Linear(dims, dims, bias\=False)
self.key\_proj \= nn.Linear(dims, dims, bias\=False)
self.value\_proj \= nn.Linear(dims, dims, bias\=False)
self.out\_proj \= nn.Linear(dims, dims, bias\=False)
def \_\_call\_\_(self, queries, keys, values, mask\=None, cache\=None):
queries \= self.query\_proj(queries)
keys \= self.key\_proj(keys)
values \= self.value\_proj(values)
\# Extract some shapes
num\_heads \= self.num\_heads
B, L, D \= queries.shape
\# Prepare the queries, keys and values for the attention computation
queries \= queries.reshape(B, L, num\_heads, \-1).transpose(0, 2, 1, 3)
keys \= keys.reshape(B, L, num\_heads, \-1).transpose(0, 2, 1, 3)
values \= values.reshape(B, L, num\_heads, \-1).transpose(0, 2, 1, 3)
\# Add RoPE to the queries and keys and combine them with the cache
if cache is not None:
key\_cache, value\_cache \= cache
queries \= self.rope(queries, offset\=key\_cache.shape\[2\])
keys \= self.rope(keys, offset\=key\_cache.shape\[2\])
keys \= mx.concatenate(\[key\_cache, keys\], axis\=2)
values \= mx.concatenate(\[value\_cache, values\], axis\=2)
else:
queries \= self.rope(queries)
keys \= self.rope(keys)
\# Finally perform the attention computation
scale \= math.sqrt(1 / queries.shape\[\-1\])
scores \= (queries \* scale) @ keys.transpose(0, 1, 3, 2)
if mask is not None:
scores \= scores + mask
scores \= mx.softmax(scores, axis\=-1)
values\_hat \= (scores @ values).transpose(0, 2, 1, 3).reshape(B, L, \-1)
\# Note that we return the keys and values to possibly be used as a cache
return self.out\_proj(values\_hat), (keys, values)
### Encoder layer[#](#encoder-layer "Link to this heading")
The other component of the Llama model is the encoder layer which uses RMS normalization [\[2\]](#id5)
and SwiGLU. [\[3\]](#id6)
For RMS normalization we will use [`mlx.nn.RMSNorm`](../python/nn/_autosummary/mlx.nn.RMSNorm.html#mlx.nn.RMSNorm "mlx.nn.RMSNorm")
that is already provided in `mlx.nn`.
class LlamaEncoderLayer(nn.Module):
def \_\_init\_\_(self, dims: int, mlp\_dims: int, num\_heads: int):
super().\_\_init\_\_()
self.attention \= LlamaAttention(dims, num\_heads)
self.norm1 \= nn.RMSNorm(dims)
self.norm2 \= nn.RMSNorm(dims)
self.linear1 \= nn.Linear(dims, mlp\_dims, bias\=False)
self.linear2 \= nn.Linear(dims, mlp\_dims, bias\=False)
self.linear3 \= nn.Linear(mlp\_dims, dims, bias\=False)
def \_\_call\_\_(self, x, mask\=None, cache\=None):
y \= self.norm1(x)
y, cache \= self.attention(y, y, y, mask, cache)
x \= x + y
y \= self.norm2(x)
a \= self.linear1(y)
b \= self.linear2(y)
y \= a \* mx.sigmoid(a) \* b
y \= self.linear3(y)
x \= x + y
return x, cache
### Full model[#](#full-model "Link to this heading")
To implement any Llama model we simply have to combine `LlamaEncoderLayer` instances with an [`mlx.nn.Embedding`](../python/nn/_autosummary/mlx.nn.Embedding.html#mlx.nn.Embedding "mlx.nn.Embedding")
to embed the input tokens.
class Llama(nn.Module):
def \_\_init\_\_(
self, num\_layers: int, vocab\_size: int, dims: int, mlp\_dims: int, num\_heads: int
):
super().\_\_init\_\_()
self.embedding \= nn.Embedding(vocab\_size, dims)
self.layers \= \[\
LlamaEncoderLayer(dims, mlp\_dims, num\_heads) for \_ in range(num\_layers)\
\]
self.norm \= nn.RMSNorm(dims)
self.out\_proj \= nn.Linear(dims, vocab\_size, bias\=False)
def \_\_call\_\_(self, x):
mask \= nn.MultiHeadAttention.create\_additive\_causal\_mask(x.shape\[1\])
mask \= mask.astype(self.embedding.weight.dtype)
x \= self.embedding(x)
for l in self.layers:
x, \_ \= l(x, mask)
x \= self.norm(x)
return self.out\_proj(x)
Note that in the implementation above we use a simple list to hold the encoder layers but using `model.parameters()` will still consider these layers.
### Generation[#](#generation "Link to this heading")
Our `Llama` module can be used for training but not inference as the `__call__` method above processes one input, completely ignores the cache and performs no sampling whatsoever. In the rest of this subsection, we will implement the inference function as a python generator that processes the prompt and then autoregressively yields tokens one at a time.
class Llama(nn.Module):
...
def generate(self, x, temp\=1.0):
cache \= \[\]
\# Make an additive causal mask. We will need that to process the prompt.
mask \= nn.MultiHeadAttention.create\_additive\_causal\_mask(x.shape\[1\])
mask \= mask.astype(self.embedding.weight.dtype)
\# First we process the prompt x the same way as in \_\_call\_\_ but
\# save the caches in cache
x \= self.embedding(x)
for l in self.layers:
x, c \= l(x, mask\=mask)
cache.append(c) \# <--- we store the per layer cache in a
\# simple python list
x \= self.norm(x)
y \= self.out\_proj(x\[:, \-1\]) \# <--- we only care about the last logits
\# that generate the next token
y \= mx.random.categorical(y \* (1/temp))
\# y now has size \[1\]
\# Since MLX is lazily evaluated nothing is computed yet.
\# Calling y.item() would force the computation to happen at
\# this point but we can also choose not to do that and let the
\# user choose when to start the computation.
yield y
\# Now we parsed the prompt and generated the first token we
\# need to feed it back into the model and loop to generate the
\# rest.
while True:
\# Unsqueezing the last dimension to add a sequence length
\# dimension of 1
x \= y\[:, None\]
x \= self.embedding(x)
for i in range(len(cache)):
\# We are overwriting the arrays in the cache list. When
\# the computation will happen, MLX will be discarding the
\# old cache the moment it is not needed anymore.
x, cache\[i\] \= self.layers\[i\](x, mask\=None, cache\=cache\[i\])
x \= self.norm(x)
y \= self.out\_proj(x\[:, \-1\])
y \= mx.random.categorical(y \* (1/temp))
yield y
### Putting it all together[#](#putting-it-all-together "Link to this heading")
We now have everything we need to create a Llama model and sample tokens from it. In the following code, we randomly initialize a small Llama model, process 6 tokens of prompt and generate 10 tokens.
model \= Llama(num\_layers\=12, vocab\_size\=8192, dims\=512, mlp\_dims\=1024, num\_heads\=8)
\# Since MLX is lazily evaluated nothing has actually been materialized yet.
\# We could have set the \`dims\` to 20\_000 on a machine with 8GB of RAM and the
\# code above would still run. Let's actually materialize the model.
mx.eval(model.parameters())
prompt \= mx.array(\[\[1, 10, 8, 32, 44, 7\]\]) \# <-- Note the double brackets because we
\# have a batch dimension even
\# though it is 1 in this case
generated \= \[t for i, t in zip(range(10), model.generate(prompt, 0.8))\]
\# Since we haven't evaluated anything, nothing is computed yet. The list
\# \`generated\` contains the arrays that hold the computation graph for the
\# full processing of the prompt and the generation of 10 tokens.
#
\# We can evaluate them one at a time, or all together. Concatenate them or
\# print them. They would all result in very similar runtimes and give exactly
\# the same results.
mx.eval(generated)
Converting the weights[#](#converting-the-weights "Link to this heading")
--------------------------------------------------------------------------
This section assumes that you have access to the original Llama weights and the SentencePiece model that comes with them. We will write a small script to convert the PyTorch weights to MLX compatible ones and write them in a NPZ file that can be loaded directly by MLX.
import argparse
from itertools import starmap
import numpy as np
import torch
def map\_torch\_to\_mlx(key, value):
if "tok\_embedding" in key:
key \= "embedding.weight"
elif "norm" in key:
key \= key.replace("attention\_norm", "norm1").replace("ffn\_norm", "norm2")
elif "wq" in key or "wk" in key or "wv" in key or "wo" in key:
key \= key.replace("wq", "query\_proj")
key \= key.replace("wk", "key\_proj")
key \= key.replace("wv", "value\_proj")
key \= key.replace("wo", "out\_proj")
elif "w1" in key or "w2" in key or "w3" in key:
\# The FFN is a separate submodule in PyTorch
key \= key.replace("feed\_forward.w1", "linear1")
key \= key.replace("feed\_forward.w3", "linear2")
key \= key.replace("feed\_forward.w2", "linear3")
elif "output" in key:
key \= key.replace("output", "out\_proj")
elif "rope" in key:
return None, None
return key, value.numpy()
if \_\_name\_\_ \== "\_\_main\_\_":
parser \= argparse.ArgumentParser(description\="Convert Llama weights to MLX")
parser.add\_argument("torch\_weights")
parser.add\_argument("output\_file")
args \= parser.parse\_args()
state \= torch.load(args.torch\_weights)
np.savez(
args.output\_file,
\*\*{k: v for k, v in starmap(map\_torch\_to\_mlx, state.items()) if k is not None}
)
Weight loading and benchmarking[#](#weight-loading-and-benchmarking "Link to this heading")
--------------------------------------------------------------------------------------------
After converting the weights to be compatible to our implementation, all that is left is to load them from disk and we can finally use the LLM to generate text. We can load numpy format files using the [`mlx.core.load()`](../python/_autosummary/mlx.core.load.html#mlx.core.load "mlx.core.load")
operation.
To create a parameter dictionary from the key/value representation of NPZ files we will use the [`mlx.utils.tree_unflatten()`](../python/_autosummary/mlx.utils.tree_unflatten.html#mlx.utils.tree_unflatten "mlx.utils.tree_unflatten")
helper method as follows:
from mlx.utils import tree\_unflatten
model.update(tree\_unflatten(list(mx.load(weight\_file).items())))
[`mlx.utils.tree_unflatten()`](../python/_autosummary/mlx.utils.tree_unflatten.html#mlx.utils.tree_unflatten "mlx.utils.tree_unflatten")
will take keys from the NPZ file that look like `layers.2.attention.query_proj.weight` and will transform them to
{"layers": \[..., ..., {"attention": {"query\_proj": {"weight": ...}}}\]}
which can then be used to update the model. Note that the method above incurs several unnecessary copies from disk to numpy and then from numpy to MLX. It will be replaced in the future with direct loading to MLX.
You can download the full example code in [mlx-examples](https://github.com/ml-explore/mlx-examples/tree/main/llms/llama)
. Assuming, the existence of `weights.pth` and `tokenizer.model` in the current working directory we can play around with our inference script as follows (the timings are representative of an M1 Ultra and the 7B parameter Llama model):
$ python convert.py weights.pth llama-7B.mlx.npz
$ python llama.py llama-7B.mlx.npz tokenizer.model 'Call me Ishmael. Some years ago never mind how long precisely'
\[INFO\] Loading model from disk: 5.247 s
Press enter to start generation
------
, having little or no money in my purse, and nothing of greater consequence in my mind, I happened to be walking down Gower Street in the afternoon, in the heavy rain, and I saw a few steps off, a man in rags, who sat upon his bundle and looked hard into the wet as if he were going to cry. I watched him attentively for some time, and could not but observe that, though a numerous crowd was hurrying up and down,
------
\[INFO\] Prompt processing: 0.437 s
\[INFO\] Full generation: 4.330 s
We observe that 4.3 seconds are required to generate 100 tokens and 0.4 seconds of those are spent processing the prompt. This amounts to a little over **39 ms per token**.
By running with a much bigger prompt we can see that the per token generation time as well as the prompt processing time remains almost constant.
$ python llama.py llama-7B.mlx.npz tokenizer.model 'Call me Ishmael. Some years ago never mind how long precisely, having little or no money in my purse, and nothing of greater consequence in my mind, I happened to be walking down Gower Street in the afternoon, in the heavy rain, and I saw a few steps off, a man in rags, who sat upon his bundle and looked hard into the wet as if he were going to cry. I watched him attentively for some time, and could not but observe that, though a numerous crowd was hurrying up and down, nobody took the least notice of him. I stopped at last, at a little distance, as if I had been in doubt, and after looking on a few minutes, walked straight up to him. He slowly raised his eyes, and fixed them upon me for a moment, without speaking, and then resumed his place and posture as before. I stood looking at him for a while, feeling very much pain at heart, and then said to him, “What are you doing there?” Something like a smile passed over his face, as he said slowly, “I am waiting for someone; but it has been three quarters of an hour now, and he has not come.” “What is it you are waiting for?” said I. Still he made no immediate reply, but again put his face down upon his hands, and did not'
\[INFO\] Loading model from disk: 5.247 s
Press enter to start generation
------
take his eyes from the ground. “What is it you are waiting for?” said I. “I am not accustomed to be thus questioned,” said he. “You look like a reasonable man—tell me, then, what are you waiting for?” “You would not understand,” he replied; “and how could you help me, if I were to tell you?” “I should not only understand, but would do all that I could,” said I. He did not
------
\[INFO\] Prompt processing: 0.579 s
\[INFO\] Full generation: 4.690 s
$ python llama.py \--num-tokens 500 llama-7B.mlx.npz tokenizer.model 'Call me Ishmael. Some years ago never mind how long precisely, having little or no money in my purse, and nothing of greater consequence in my mind, I happened to be walking down Gower Street in the afternoon, in the heavy rain, and I saw a few steps off, a man in rags, who sat upon his bundle and looked hard into the wet as if he were going to cry. I watched him attentively for some time, and could not but observe that, though a numerous crowd was hurrying up and down, nobody took the least notice of him. I stopped at last, at a little distance, as if I had been in doubt, and after looking on a few minutes, walked straight up to him. He slowly raised his eyes, and fixed them upon me for a moment, without speaking, and then resumed his place and posture as before. I stood looking at him for a while, feeling very much pain at heart, and then said to him, “What are you doing there?” Something like a smile passed over his face, as he said slowly, “I am waiting for someone; but it has been three quarters of an hour now, and he has not come.” “What is it you are waiting for?” said I. Still he made no immediate reply, but again put his face down upon his hands, and did not'
\[INFO\] Loading model from disk: 5.628 s
Press enter to start generation
------
take his eyes from the ground. “What is it you are waiting for?” said I. “I am not accustomed to be thus questioned,” said he. “You look like a reasonable man—tell me, then, what are you waiting for?” “You would not understand,” he replied; “and how could you help me, if I were to tell you?” “I should not only understand, but would do all that I could,” said I. He did not reply, but still went on looking at the ground, and took hold of his bundle with a nervous trembling. I waited some time, and then resumed. “It is of no use to say you would not understand, if I were to tell you,” said he. “I have not told you why I am waiting for him,” said I. “And I am sure I should not understand,” replied he. “I will tell you then,” said I, “and, perhaps, you would not be surprised.” “No matter,” said he, “I shall be surprised anyhow; so tell me why you are waiting for him.” “He is my friend,” said I. “Yes,” said he, with a slight smile, “I know.” “He has been kind to me,” said I, “and I am waiting for him. I want to see him, and could have waited as I am now, for a much longer time.” “He will not soon come,” said he. “Unless he sees you here, he will not know of your having waited, and he will be very unlikely to come.” “No matter,” said I, “I shall wait for him.” “This is a strange thing,” said he, still with the same amused smile. “How did you know,” said I, “that he was coming? How should you be waiting?” “That is my secret,” said he. “And you expect him?” “Yes,” said I. “Are you disappointed then, if he does not come?” “No,” said I, “it is his secret, not mine.” “If he comes,” said he, “do you mean to go straight away?” “Yes,” said I, “I cannot be happy if I do not go straight away after him.” “Did you know this place before?” asked he. “Yes,” said I. “Is there any shop to buy food here?” “
------
\[INFO\] Prompt processing: 0.633 s
\[INFO\] Full generation: 21.475 s
Scripts[#](#scripts "Link to this heading")
--------------------------------------------
Download the code
The full example code is available in [mlx-examples](https://github.com/ml-explore/mlx-examples/tree/main/llms/llama)
.
Contents
---
# Distributed Communication — MLX 0.24.0 documentation
[Skip to main content](#main-content)
Back to top
Ctrl+K
[](../index.html)
[](https://github.com/ml-explore/mlx "Source repository")
* [.rst](../_sources/usage/distributed.rst "Download source file")
* .pdf
Distributed Communication
=========================
Contents
--------
Distributed Communication[#](#distributed-communication "Link to this heading")
================================================================================
MLX supports distributed communication operations that allow the computational cost of training or inference to be shared across many physical machines. At the moment we support two different communication backends:
* [MPI](https://en.wikipedia.org/wiki/Message_Passing_Interface)
a full-featured and mature distributed communications library
* A **ring** backend of our own that uses native TCP sockets and should be faster for thunderbolt connections.
The list of all currently supported operations and their documentation can be seen in the [API docs](../python/distributed.html#distributed)
.
Note
Some operations may not be supported or not as fast as they should be. We are adding more and tuning the ones we have as we are figuring out the best way to do distributed computing on Macs using MLX.
Getting Started[#](#getting-started "Link to this heading")
------------------------------------------------------------
A distributed program in MLX is as simple as:
import mlx.core as mx
world \= mx.distributed.init()
x \= mx.distributed.all\_sum(mx.ones(10))
print(world.rank(), x)
The program above sums the array `mx.ones(10)` across all distributed processes. However, when this script is run with `python` only one process is launched and no distributed communication takes place. Namely, all operations in `mx.distributed` are noops when the distributed group has a size of one. This property allows us to avoid code that checks if we are in a distributed setting similar to the one below:
import mlx.core as mx
x \= ...
world \= mx.distributed.init()
\# No need for the check we can simply do x = mx.distributed.all\_sum(x)
if world.size() \> 1:
x \= mx.distributed.all\_sum(x)
### Running Distributed Programs[#](#running-distributed-programs "Link to this heading")
MLX provides `mlx.launch` a helper script to launch distributed programs. Continuing with our initial example we can run it on localhost with 4 processes using
$ mlx.launch \-n 4 my\_script.py
3 array(\[4, 4, 4, ..., 4, 4, 4\], dtype\=float32)
2 array(\[4, 4, 4, ..., 4, 4, 4\], dtype\=float32)
1 array(\[4, 4, 4, ..., 4, 4, 4\], dtype\=float32)
0 array(\[4, 4, 4, ..., 4, 4, 4\], dtype\=float32)
We can also run it on some remote hosts by providing their IPs (provided that the script exists on all hosts and they are reachable by ssh)
$ mlx.launch \--hosts ip1,ip2,ip3,ip4 my\_script.py
3 array(\[4, 4, 4, ..., 4, 4, 4\], dtype\=float32)
2 array(\[4, 4, 4, ..., 4, 4, 4\], dtype\=float32)
1 array(\[4, 4, 4, ..., 4, 4, 4\], dtype\=float32)
0 array(\[4, 4, 4, ..., 4, 4, 4\], dtype\=float32)
Consult the dedicated [usage guide](launching_distributed.html)
for more information on using `mlx.launch`.
### Selecting Backend[#](#selecting-backend "Link to this heading")
You can select the backend you want to use when calling [`init()`](../python/_autosummary/mlx.core.distributed.init.html#mlx.core.distributed.init "mlx.core.distributed.init")
by passing one of `{'any', 'ring', 'mpi'}`. When passing `any`, MLX will try to initialize the `ring` backend and if it fails the `mpi` backend. If they both fail then a singleton group is created.
Note
After a distributed backend is successfully initialized [`init()`](../python/_autosummary/mlx.core.distributed.init.html#mlx.core.distributed.init "mlx.core.distributed.init")
will return **the same backend** if called without arguments or with backend set to `any`.
The following examples aim to clarify the backend initialization logic in MLX:
\# Case 1: Initialize MPI regardless if it was possible to initialize the ring backend
world \= mx.distributed.init(backend\="mpi")
world2 \= mx.distributed.init() \# subsequent calls return the MPI backend!
\# Case 2: Initialize any backend
world \= mx.distributed.init(backend\="any") \# equivalent to no arguments
world2 \= mx.distributed.init() \# same as above
\# Case 3: Initialize both backends at the same time
world\_mpi \= mx.distributed.init(backend\="mpi")
world\_ring \= mx.distributed.init(backend\="ring")
world\_any \= mx.distributed.init() \# same as MPI because it was initialized first!
Training Example[#](#training-example "Link to this heading")
--------------------------------------------------------------
In this section we will adapt an MLX training loop to support data parallel distributed training. Namely, we will average the gradients across a set of hosts before applying them to the model.
Our training loop looks like the following code snippet if we omit the model, dataset and optimizer initialization.
model \= ...
optimizer \= ...
dataset \= ...
def step(model, x, y):
loss, grads \= loss\_grad\_fn(model, x, y)
optimizer.update(model, grads)
return loss
for x, y in dataset:
loss \= step(model, x, y)
mx.eval(loss, model.parameters())
All we have to do to average the gradients across machines is perform an [`all_sum()`](../python/_autosummary/mlx.core.distributed.all_sum.html#mlx.core.distributed.all_sum "mlx.core.distributed.all_sum")
and divide by the size of the [`Group`](../python/_autosummary/mlx.core.distributed.Group.html#mlx.core.distributed.Group "mlx.core.distributed.Group")
. Namely we have to [`mlx.utils.tree_map()`](../python/_autosummary/mlx.utils.tree_map.html#mlx.utils.tree_map "mlx.utils.tree_map")
the gradients with following function.
def all\_avg(x):
return mx.distributed.all\_sum(x) / mx.distributed.init().size()
Putting everything together our training loop step looks as follows with everything else remaining the same.
from mlx.utils import tree\_map
def all\_reduce\_grads(grads):
N \= mx.distributed.init().size()
if N \== 1:
return grads
return tree\_map(
lambda x: mx.distributed.all\_sum(x) / N,
grads
)
def step(model, x, y):
loss, grads \= loss\_grad\_fn(model, x, y)
grads \= all\_reduce\_grads(grads) \# <--- This line was added
optimizer.update(model, grads)
return loss
### Utilizing `nn.average_gradients`[#](#utilizing-nn-average-gradients "Link to this heading")
Although the code example above works correctly; it performs one communication per gradient. It is significantly more efficient to aggregate several gradients together and perform fewer communication steps.
This is the purpose of [`mlx.nn.average_gradients()`](../python/_autosummary/mlx.nn.average_gradients.html#mlx.nn.average_gradients "mlx.nn.average_gradients")
. The final code looks almost identical to the example above:
model \= ...
optimizer \= ...
dataset \= ...
def step(model, x, y):
loss, grads \= loss\_grad\_fn(model, x, y)
grads \= mlx.nn.average\_gradients(grads) \# <---- This line was added
optimizer.update(model, grads)
return loss
for x, y in dataset:
loss \= step(model, x, y)
mx.eval(loss, model.parameters())
Getting Started with MPI[#](#getting-started-with-mpi "Link to this heading")
------------------------------------------------------------------------------
MLX already comes with the ability to “talk” to MPI if it is installed on the machine. Launching distributed MLX programs that use MPI can be done with `mpirun` as expected. However, in the following examples we will be using `mlx.launch --backend mpi` which takes care of some nuisances such as setting absolute paths for the `mpirun` executable and the `libmpi.dyld` shared library.
The simplest possible usage is the following which, assuming the minimal example in the beginning of this page, should result in:
$ mlx.launch \--backend mpi \-n 2 test.py
1 array(\[2, 2, 2, ..., 2, 2, 2\], dtype\=float32)
0 array(\[2, 2, 2, ..., 2, 2, 2\], dtype\=float32)
The above launches two processes on the same (local) machine and we can see both standard output streams. The processes send the array of 1s to each other and compute the sum which is printed. Launching with `mlx.launch -n 4 ...` would print 4 etc.
### Installing MPI[#](#installing-mpi "Link to this heading")
MPI can be installed with Homebrew, using the Anaconda package manager or compiled from source. Most of our testing is done using `openmpi` installed with the Anaconda package manager as follows:
$ conda install conda-forge::openmpi
Installing with Homebrew may require specifying the location of `libmpi.dyld` so that MLX can find it and load it at runtime. This can simply be achieved by passing the `DYLD_LIBRARY_PATH` environment variable to `mpirun` and it is done automatically by `mlx.launch`.
$ mpirun \-np 2 \-x DYLD\_LIBRARY\_PATH\=/opt/homebrew/lib/ python test.py
$ \# or simply
$ mlx.launch \-n 2 test.py
### Setting up Remote Hosts[#](#setting-up-remote-hosts "Link to this heading")
MPI can automatically connect to remote hosts and set up the communication over the network if the remote hosts can be accessed via ssh. A good checklist to debug connectivity issues is the following:
* `ssh hostname` works from all machines to all machines without asking for password or host confirmation
* `mpirun` is accessible on all machines.
* Ensure that the `hostname` used by MPI is the one that you have configured in the `.ssh/config` files on all machines.
### Tuning MPI All Reduce[#](#tuning-mpi-all-reduce "Link to this heading")
Note
For faster all reduce consider using the ring backend either with Thunderbolt connections or over Ethernet.
Configure MPI to use N tcp connections between each host to improve bandwidth by passing `--mca btl_tcp_links N`.
Force MPI to use the most performant network interface by setting `--mca btl_tcp_if_include ` where `` should be the interface you want to use.
Getting Started with Ring[#](#getting-started-with-ring "Link to this heading")
--------------------------------------------------------------------------------
The ring backend does not depend on any third party library so it is always available. It uses TCP sockets so the nodes need to be reachable via a network. As the name suggests the nodes are connected in a ring which means that rank 1 can only communicate with rank 0 and rank 2, rank 2 only with rank 1 and rank 3 and so on and so forth. As a result [`send()`](../python/_autosummary/mlx.core.distributed.send.html#mlx.core.distributed.send "mlx.core.distributed.send")
and [`recv()`](../python/_autosummary/mlx.core.distributed.recv.html#mlx.core.distributed.recv "mlx.core.distributed.recv")
with arbitrary sender and receiver is not supported in the ring backend.
### Defining a Ring[#](#defining-a-ring "Link to this heading")
The easiest way to define and use a ring is via a JSON hostfile and the `mlx.launch` [helper script](launching_distributed.html)
. For each node one defines a hostname to ssh into to run commands on this node and one or more IPs that this node will listen to for connections.
For example the hostfile below defines a 4 node ring. `hostname1` will be rank 0, `hostname2` rank 1 etc.
\[\
{"ssh": "hostname1", "ips": \["123.123.123.1"\]},\
{"ssh": "hostname2", "ips": \["123.123.123.2"\]},\
{"ssh": "hostname3", "ips": \["123.123.123.3"\]},\
{"ssh": "hostname4", "ips": \["123.123.123.4"\]}\
\]
Running `mlx.launch --hostfile ring-4.json my_script.py` will ssh into each node, run the script which will listen for connections in each of the provided IPs. Specifically, `hostname1` will connect to `123.123.123.2` and accept a connection from `123.123.123.4` and so on and so forth.
### Thunderbolt Ring[#](#thunderbolt-ring "Link to this heading")
Although the ring backend can have benefits over MPI even for Ethernet, its main purpose is to use Thunderbolt rings for higher bandwidth communication. Setting up such thunderbolt rings can be done manually, but is a relatively tedious process. To simplify this, we provide the utility `mlx.distributed_config`.
To use `mlx.distributed_config` your computers need to be accessible by ssh via Ethernet or Wi-Fi. Subsequently, connect them via thunderbolt cables and then call the utility as follows:
mlx.distributed\_config \--verbose \--hosts host1,host2,host3,host4
By default the script will attempt to discover the thunderbolt ring and provide you with the commands to configure each node as well as the `hostfile.json` to use with `mlx.launch`. If password-less `sudo` is available on the nodes then `--auto-setup` can be used to configure them automatically.
To validate your connection without configuring anything `mlx.distributed_config` can also plot the ring using DOT format.
mlx.distributed\_config \--verbose \--hosts host1,host2,host3,host4 \--dot \>ring.dot
dot \-Tpng ring.dot \>ring.png
open ring.png
If you want to go through the process manually, the steps are as follows:
* Disable the thunderbolt bridge interface
* For the cable connecting rank `i` to rank `i + 1` find the interfaces corresponding to that cable in nodes `i` and `i + 1`.
* Set up a unique subnetwork connecting the two nodes for the corresponding interfaces. For instance if the cable corresponds to `en2` on node `i` and `en2` also on node `i + 1` then we may assign IPs `192.168.0.1` and `192.168.0.2` respectively to the two nodes. For more details you can see the commands prepared by the utility script.
Contents
---
# Array — MLX 0.24.0 documentation
[Skip to main content](#main-content)
Back to top
Ctrl+K
[](../index.html)
[](https://github.com/ml-explore/mlx "Source repository")
* [.rst](../_sources/python/array.rst "Download source file")
* .pdf
Array
=====
Array[#](#array "Link to this heading")
========================================
| | |
| --- | --- |
| [`array`](_autosummary/mlx.core.array.html#mlx.core.array "mlx.core.array") | An N-dimensional array object. |
| [`array.astype`](_autosummary/mlx.core.array.astype.html#mlx.core.array.astype "mlx.core.array.astype")
(self, dtype\[, stream\]) | Cast the array to a specified type. |
| [`array.at`](_autosummary/mlx.core.array.at.html#mlx.core.array.at "mlx.core.array.at") | Used to apply updates at the given indices. |
| [`array.item`](_autosummary/mlx.core.array.item.html#mlx.core.array.item "mlx.core.array.item")
(self) | Access the value of a scalar array. |
| [`array.tolist`](_autosummary/mlx.core.array.tolist.html#mlx.core.array.tolist "mlx.core.array.tolist")
(self) | Convert the array to a Python [`list`](https://docs.python.org/3/library/stdtypes.html#list "(in Python v3.13)")
. |
| [`array.dtype`](_autosummary/mlx.core.array.dtype.html#mlx.core.array.dtype "mlx.core.array.dtype") | The array's [`Dtype`](_autosummary/mlx.core.Dtype.html#mlx.core.Dtype "mlx.core.Dtype")
. |
| [`array.itemsize`](_autosummary/mlx.core.array.itemsize.html#mlx.core.array.itemsize "mlx.core.array.itemsize") | The size of the array's datatype in bytes. |
| [`array.nbytes`](_autosummary/mlx.core.array.nbytes.html#mlx.core.array.nbytes "mlx.core.array.nbytes") | The number of bytes in the array. |
| [`array.ndim`](_autosummary/mlx.core.array.ndim.html#mlx.core.array.ndim "mlx.core.array.ndim") | The array's dimension. |
| [`array.shape`](_autosummary/mlx.core.array.shape.html#mlx.core.array.shape "mlx.core.array.shape") | The shape of the array as a Python tuple. |
| [`array.size`](_autosummary/mlx.core.array.size.html#mlx.core.array.size "mlx.core.array.size") | Number of elements in the array. |
| [`array.abs`](_autosummary/mlx.core.array.abs.html#mlx.core.array.abs "mlx.core.array.abs")
(self, \*\[, stream\]) | See [`abs()`](_autosummary/mlx.core.abs.html#mlx.core.abs "mlx.core.abs")
. |
| [`array.all`](_autosummary/mlx.core.array.all.html#mlx.core.array.all "mlx.core.array.all")
(self\[, axis, keepdims, stream\]) | See [`all()`](_autosummary/mlx.core.all.html#mlx.core.all "mlx.core.all")
. |
| [`array.any`](_autosummary/mlx.core.array.any.html#mlx.core.array.any "mlx.core.array.any")
(self\[, axis, keepdims, stream\]) | See [`any()`](_autosummary/mlx.core.any.html#mlx.core.any "mlx.core.any")
. |
| [`array.argmax`](_autosummary/mlx.core.array.argmax.html#mlx.core.array.argmax "mlx.core.array.argmax")
(self\[, axis, keepdims, stream\]) | See [`argmax()`](_autosummary/mlx.core.argmax.html#mlx.core.argmax "mlx.core.argmax")
. |
| [`array.argmin`](_autosummary/mlx.core.array.argmin.html#mlx.core.array.argmin "mlx.core.array.argmin")
(self\[, axis, keepdims, stream\]) | See [`argmin()`](_autosummary/mlx.core.argmin.html#mlx.core.argmin "mlx.core.argmin")
. |
| [`array.conj`](_autosummary/mlx.core.array.conj.html#mlx.core.array.conj "mlx.core.array.conj")
(self, \*\[, stream\]) | See [`conj()`](_autosummary/mlx.core.conj.html#mlx.core.conj "mlx.core.conj")
. |
| [`array.cos`](_autosummary/mlx.core.array.cos.html#mlx.core.array.cos "mlx.core.array.cos")
(self, \*\[, stream\]) | See [`cos()`](_autosummary/mlx.core.cos.html#mlx.core.cos "mlx.core.cos")
. |
| [`array.cummax`](_autosummary/mlx.core.array.cummax.html#mlx.core.array.cummax "mlx.core.array.cummax")
(self\[, axis, reverse, ...\]) | See [`cummax()`](_autosummary/mlx.core.cummax.html#mlx.core.cummax "mlx.core.cummax")
. |
| [`array.cummin`](_autosummary/mlx.core.array.cummin.html#mlx.core.array.cummin "mlx.core.array.cummin")
(self\[, axis, reverse, ...\]) | See [`cummin()`](_autosummary/mlx.core.cummin.html#mlx.core.cummin "mlx.core.cummin")
. |
| [`array.cumprod`](_autosummary/mlx.core.array.cumprod.html#mlx.core.array.cumprod "mlx.core.array.cumprod")
(self\[, axis, reverse, ...\]) | See [`cumprod()`](_autosummary/mlx.core.cumprod.html#mlx.core.cumprod "mlx.core.cumprod")
. |
| [`array.cumsum`](_autosummary/mlx.core.array.cumsum.html#mlx.core.array.cumsum "mlx.core.array.cumsum")
(self\[, axis, reverse, ...\]) | See [`cumsum()`](_autosummary/mlx.core.cumsum.html#mlx.core.cumsum "mlx.core.cumsum")
. |
| [`array.diag`](_autosummary/mlx.core.array.diag.html#mlx.core.array.diag "mlx.core.array.diag")
(self\[, k, stream\]) | Extract a diagonal or construct a diagonal matrix. |
| [`array.diagonal`](_autosummary/mlx.core.array.diagonal.html#mlx.core.array.diagonal "mlx.core.array.diagonal")
(self\[, offset, axis1, axis2, ...\]) | See [`diagonal()`](_autosummary/mlx.core.diagonal.html#mlx.core.diagonal "mlx.core.diagonal")
. |
| [`array.exp`](_autosummary/mlx.core.array.exp.html#mlx.core.array.exp "mlx.core.array.exp")
(self, \*\[, stream\]) | See [`exp()`](_autosummary/mlx.core.exp.html#mlx.core.exp "mlx.core.exp")
. |
| [`array.flatten`](_autosummary/mlx.core.array.flatten.html#mlx.core.array.flatten "mlx.core.array.flatten")
(self\[, start\_axis, end\_axis, ...\]) | See [`flatten()`](_autosummary/mlx.core.flatten.html#mlx.core.flatten "mlx.core.flatten")
. |
| [`array.log`](_autosummary/mlx.core.array.log.html#mlx.core.array.log "mlx.core.array.log")
(self, \*\[, stream\]) | See [`log()`](_autosummary/mlx.core.log.html#mlx.core.log "mlx.core.log")
. |
| [`array.log10`](_autosummary/mlx.core.array.log10.html#mlx.core.array.log10 "mlx.core.array.log10")
(self, \*\[, stream\]) | See [`log10()`](_autosummary/mlx.core.log10.html#mlx.core.log10 "mlx.core.log10")
. |
| [`array.log1p`](_autosummary/mlx.core.array.log1p.html#mlx.core.array.log1p "mlx.core.array.log1p")
(self, \*\[, stream\]) | See [`log1p()`](_autosummary/mlx.core.log1p.html#mlx.core.log1p "mlx.core.log1p")
. |
| [`array.log2`](_autosummary/mlx.core.array.log2.html#mlx.core.array.log2 "mlx.core.array.log2")
(self, \*\[, stream\]) | See [`log2()`](_autosummary/mlx.core.log2.html#mlx.core.log2 "mlx.core.log2")
. |
| [`array.logsumexp`](_autosummary/mlx.core.array.logsumexp.html#mlx.core.array.logsumexp "mlx.core.array.logsumexp")
(self\[, axis, keepdims, stream\]) | See [`logsumexp()`](_autosummary/mlx.core.logsumexp.html#mlx.core.logsumexp "mlx.core.logsumexp")
. |
| [`array.max`](_autosummary/mlx.core.array.max.html#mlx.core.array.max "mlx.core.array.max")
(self\[, axis, keepdims, stream\]) | See [`max()`](_autosummary/mlx.core.max.html#mlx.core.max "mlx.core.max")
. |
| [`array.mean`](_autosummary/mlx.core.array.mean.html#mlx.core.array.mean "mlx.core.array.mean")
(self\[, axis, keepdims, stream\]) | See [`mean()`](_autosummary/mlx.core.mean.html#mlx.core.mean "mlx.core.mean")
. |
| [`array.min`](_autosummary/mlx.core.array.min.html#mlx.core.array.min "mlx.core.array.min")
(self\[, axis, keepdims, stream\]) | See [`min()`](_autosummary/mlx.core.min.html#mlx.core.min "mlx.core.min")
. |
| [`array.moveaxis`](_autosummary/mlx.core.array.moveaxis.html#mlx.core.array.moveaxis "mlx.core.array.moveaxis")
(self, source, destination, \*) | See [`moveaxis()`](_autosummary/mlx.core.moveaxis.html#mlx.core.moveaxis "mlx.core.moveaxis")
. |
| [`array.prod`](_autosummary/mlx.core.array.prod.html#mlx.core.array.prod "mlx.core.array.prod")
(self\[, axis, keepdims, stream\]) | See [`prod()`](_autosummary/mlx.core.prod.html#mlx.core.prod "mlx.core.prod")
. |
| [`array.reciprocal`](_autosummary/mlx.core.array.reciprocal.html#mlx.core.array.reciprocal "mlx.core.array.reciprocal")
(self, \*\[, stream\]) | See [`reciprocal()`](_autosummary/mlx.core.reciprocal.html#mlx.core.reciprocal "mlx.core.reciprocal")
. |
| [`array.reshape`](_autosummary/mlx.core.array.reshape.html#mlx.core.array.reshape "mlx.core.array.reshape")
(self, \*shape\[, stream\]) | Equivalent to [`reshape()`](_autosummary/mlx.core.reshape.html#mlx.core.reshape "mlx.core.reshape")
but the shape can be passed either as a [`tuple`](https://docs.python.org/3/library/stdtypes.html#tuple "(in Python v3.13)")
or as separate arguments. |
| [`array.round`](_autosummary/mlx.core.array.round.html#mlx.core.array.round "mlx.core.array.round")
(self\[, decimals, stream\]) | See [`round()`](_autosummary/mlx.core.round.html#mlx.core.round "mlx.core.round")
. |
| [`array.rsqrt`](_autosummary/mlx.core.array.rsqrt.html#mlx.core.array.rsqrt "mlx.core.array.rsqrt")
(self, \*\[, stream\]) | See [`rsqrt()`](_autosummary/mlx.core.rsqrt.html#mlx.core.rsqrt "mlx.core.rsqrt")
. |
| [`array.sin`](_autosummary/mlx.core.array.sin.html#mlx.core.array.sin "mlx.core.array.sin")
(self, \*\[, stream\]) | See [`sin()`](_autosummary/mlx.core.sin.html#mlx.core.sin "mlx.core.sin")
. |
| [`array.split`](_autosummary/mlx.core.array.split.html#mlx.core.array.split "mlx.core.array.split")
(self, indices\_or\_sections\[, ...\]) | See [`split()`](_autosummary/mlx.core.split.html#mlx.core.split "mlx.core.split")
. |
| [`array.sqrt`](_autosummary/mlx.core.array.sqrt.html#mlx.core.array.sqrt "mlx.core.array.sqrt")
(self, \*\[, stream\]) | See [`sqrt()`](_autosummary/mlx.core.sqrt.html#mlx.core.sqrt "mlx.core.sqrt")
. |
| [`array.square`](_autosummary/mlx.core.array.square.html#mlx.core.array.square "mlx.core.array.square")
(self, \*\[, stream\]) | See [`square()`](_autosummary/mlx.core.square.html#mlx.core.square "mlx.core.square")
. |
| [`array.squeeze`](_autosummary/mlx.core.array.squeeze.html#mlx.core.array.squeeze "mlx.core.array.squeeze")
(self\[, axis, stream\]) | See [`squeeze()`](_autosummary/mlx.core.squeeze.html#mlx.core.squeeze "mlx.core.squeeze")
. |
| [`array.std`](_autosummary/mlx.core.array.std.html#mlx.core.array.std "mlx.core.array.std")
(self\[, axis, keepdims, ddof, stream\]) | See [`std()`](_autosummary/mlx.core.std.html#mlx.core.std "mlx.core.std")
. |
| [`array.sum`](_autosummary/mlx.core.array.sum.html#mlx.core.array.sum "mlx.core.array.sum")
(self\[, axis, keepdims, stream\]) | See [`sum()`](_autosummary/mlx.core.sum.html#mlx.core.sum "mlx.core.sum")
. |
| [`array.swapaxes`](_autosummary/mlx.core.array.swapaxes.html#mlx.core.array.swapaxes "mlx.core.array.swapaxes")
(self, axis1, axis2, \*\[, stream\]) | See [`swapaxes()`](_autosummary/mlx.core.swapaxes.html#mlx.core.swapaxes "mlx.core.swapaxes")
. |
| [`array.transpose`](_autosummary/mlx.core.array.transpose.html#mlx.core.array.transpose "mlx.core.array.transpose")
(self, \*axes\[, stream\]) | Equivalent to [`transpose()`](_autosummary/mlx.core.transpose.html#mlx.core.transpose "mlx.core.transpose")
but the axes can be passed either as a tuple or as separate arguments. |
| [`array.T`](_autosummary/mlx.core.array.T.html#mlx.core.array.T "mlx.core.array.T") | Equivalent to calling `self.transpose()` with no arguments. |
| [`array.var`](_autosummary/mlx.core.array.var.html#mlx.core.array.var "mlx.core.array.var")
(self\[, axis, keepdims, ddof, stream\]) | See [`var()`](_autosummary/mlx.core.var.html#mlx.core.var "mlx.core.var")
. |
| [`array.view`](_autosummary/mlx.core.array.view.html#mlx.core.array.view "mlx.core.array.view")
(self, dtype, \*\[, stream\]) | See [`view()`](_autosummary/mlx.core.view.html#mlx.core.view "mlx.core.view")
. |
---
# Devices and Streams — MLX 0.24.0 documentation
[Skip to main content](#main-content)
Back to top
Ctrl+K
[](../index.html)
[](https://github.com/ml-explore/mlx "Source repository")
* [.rst](../_sources/python/devices_and_streams.rst "Download source file")
* .pdf
Devices and Streams
===================
Devices and Streams[#](#devices-and-streams "Link to this heading")
====================================================================
| | |
| --- | --- |
| [`Device`](_autosummary/mlx.core.Device.html#mlx.core.Device "mlx.core.Device") | A device to run operations on. |
| [`Stream`](_autosummary/stream_class.html#mlx.core.Stream "mlx.core.Stream") | A stream for running operations on a given device. |
| [`default_device`](_autosummary/mlx.core.default_device.html#mlx.core.default_device "mlx.core.default_device")
() | Get the default device. |
| [`set_default_device`](_autosummary/mlx.core.set_default_device.html#mlx.core.set_default_device "mlx.core.set_default_device")
(device) | Set the default device. |
| [`default_stream`](_autosummary/mlx.core.default_stream.html#mlx.core.default_stream "mlx.core.default_stream")
(device) | Get the device's default stream. |
| [`new_stream`](_autosummary/mlx.core.new_stream.html#mlx.core.new_stream "mlx.core.new_stream")
(device) | Make a new stream on the given device. |
| [`set_default_stream`](_autosummary/mlx.core.set_default_stream.html#mlx.core.set_default_stream "mlx.core.set_default_stream")
(stream) | Set the default stream. |
| [`stream`](_autosummary/mlx.core.stream.html#mlx.core.stream "mlx.core.stream")
(s) | Create a context manager to set the default device and stream. |
| [`synchronize`](_autosummary/mlx.core.synchronize.html#mlx.core.synchronize "mlx.core.synchronize")
(\[stream\]) | Synchronize with the given stream. |
---
# Data Types — MLX 0.24.0 documentation
[Skip to main content](#main-content)
Back to top
Ctrl+K
[](../index.html)
[](https://github.com/ml-explore/mlx "Source repository")
* [.rst](../_sources/python/data_types.rst "Download source file")
* .pdf
Data Types
==========
Data Types[#](#data-types "Link to this heading")
==================================================
The default floating point type is `float32` and the default integer type is `int32`. The table below shows supported values for [`Dtype`](_autosummary/mlx.core.Dtype.html#mlx.core.Dtype "mlx.core.Dtype")
.
| | | |
| --- | --- | --- |Supported Data Types[#](#id2 "Link to this table")
| 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` | 4 | 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 [`DtypeCategory`](_autosummary/mlx.core.DtypeCategory.html#mlx.core.DtypeCategory "mlx.core.DtypeCategory")
object documentation for more information. Use [`issubdtype()`](_autosummary/mlx.core.issubdtype.html#mlx.core.issubdtype "mlx.core.issubdtype")
to determine if one `dtype` (or category) is a subtype of another category.
| | |
| --- | --- |
| [`Dtype`](_autosummary/mlx.core.Dtype.html#mlx.core.Dtype "mlx.core.Dtype") | An object to hold the type of a [`array`](_autosummary/mlx.core.array.html#mlx.core.array "mlx.core.array")
. |
| [`DtypeCategory`](_autosummary/mlx.core.DtypeCategory.html#mlx.core.DtypeCategory "mlx.core.DtypeCategory")
(value) | Type to hold categories of [`dtypes`](_autosummary/mlx.core.Dtype.html#mlx.core.Dtype "mlx.core.Dtype")
. |
| [`issubdtype`](_autosummary/mlx.core.issubdtype.html#mlx.core.issubdtype "mlx.core.issubdtype")
(arg1, arg2) | Check if a [`Dtype`](_autosummary/mlx.core.Dtype.html#mlx.core.Dtype "mlx.core.Dtype")
or [`DtypeCategory`](_autosummary/mlx.core.DtypeCategory.html#mlx.core.DtypeCategory "mlx.core.DtypeCategory")
is a subtype of another. |
| [`finfo`](_autosummary/mlx.core.finfo.html#mlx.core.finfo "mlx.core.finfo") | Get information on floating-point types. |
---
# Export Functions — MLX 0.24.0 documentation
[Skip to main content](#main-content)
Back to top
Ctrl+K
[](../index.html)
[](https://github.com/ml-explore/mlx "Source repository")
* [.rst](../_sources/python/export.rst "Download source file")
* .pdf
Export Functions
================
Export Functions[#](#export-functions "Link to this heading")
==============================================================
| | |
| --- | --- |
| [`export_function`](_autosummary/mlx.core.export_function.html#mlx.core.export_function "mlx.core.export_function")
(file, fun, \*args\[, shapeless\]) | Export a function to a file. |
| [`import_function`](_autosummary/mlx.core.import_function.html#mlx.core.import_function "mlx.core.import_function")
(file) | Import a function from a file. |
| [`exporter`](_autosummary/mlx.core.exporter.html#mlx.core.exporter "mlx.core.exporter")
(file, fun, \*\[, shapeless\]) | Make a callable object to export multiple traces of a function to a file. |
| [`export_to_dot`](_autosummary/mlx.core.export_to_dot.html#mlx.core.export_to_dot "mlx.core.export_to_dot")
(file, \*args, \*\*kwargs) | Export a graph to DOT format for visualization. |
---
# Random — MLX 0.24.0 documentation
[Skip to main content](#main-content)
Back to top
Ctrl+K
[](../index.html)
[](https://github.com/ml-explore/mlx "Source repository")
* [.rst](../_sources/python/random.rst "Download source file")
* .pdf
Random
======
Random[#](#random "Link to this heading")
==========================================
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:
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:
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.
| | |
| --- | --- |
| [`bernoulli`](_autosummary/mlx.core.random.bernoulli.html#mlx.core.random.bernoulli "mlx.core.random.bernoulli")
(\[p, shape, key, stream\]) | Generate Bernoulli random values. |
| [`categorical`](_autosummary/mlx.core.random.categorical.html#mlx.core.random.categorical "mlx.core.random.categorical")
(logits\[, axis, shape, ...\]) | Sample from a categorical distribution. |
| [`gumbel`](_autosummary/mlx.core.random.gumbel.html#mlx.core.random.gumbel "mlx.core.random.gumbel")
(\[shape, dtype, key, stream\]) | Sample from the standard Gumbel distribution. |
| [`key`](_autosummary/mlx.core.random.key.html#mlx.core.random.key "mlx.core.random.key")
(seed) | Get a PRNG key from a seed. |
| [`normal`](_autosummary/mlx.core.random.normal.html#mlx.core.random.normal "mlx.core.random.normal")
(\[shape, dtype, loc, scale, key, stream\]) | Generate normally distributed random numbers. |
| [`multivariate_normal`](_autosummary/mlx.core.random.multivariate_normal.html#mlx.core.random.multivariate_normal "mlx.core.random.multivariate_normal")
(mean, cov\[, shape, ...\]) | Generate jointly-normal random samples given a mean and covariance. |
| [`randint`](_autosummary/mlx.core.random.randint.html#mlx.core.random.randint "mlx.core.random.randint")
(low, high\[, shape, dtype, key, stream\]) | Generate random integers from the given interval. |
| [`seed`](_autosummary/mlx.core.random.seed.html#mlx.core.random.seed "mlx.core.random.seed")
(seed) | Seed the global PRNG. |
| [`split`](_autosummary/mlx.core.random.split.html#mlx.core.random.split "mlx.core.random.split")
(key\[, num, stream\]) | Split a PRNG key into sub keys. |
| [`truncated_normal`](_autosummary/mlx.core.random.truncated_normal.html#mlx.core.random.truncated_normal "mlx.core.random.truncated_normal")
(lower, upper\[, shape, ...\]) | Generate values from a truncated normal distribution. |
| [`uniform`](_autosummary/mlx.core.random.uniform.html#mlx.core.random.uniform "mlx.core.random.uniform")
(\[low, high, shape, dtype, key, stream\]) | Generate uniformly distributed random numbers. |
| [`laplace`](_autosummary/mlx.core.random.laplace.html#mlx.core.random.laplace "mlx.core.random.laplace")
(\[shape, dtype, loc, scale, key, stream\]) | Sample numbers from a Laplace distribution. |
| [`permutation`](_autosummary/mlx.core.random.permutation.html#mlx.core.random.permutation "mlx.core.random.permutation")
(x\[, axis, key, stream\]) | Generate a random permutation or permute the entries of an array. |
---
# Operations — MLX 0.24.0 documentation
[Skip to main content](#main-content)
Back to top
Ctrl+K
[](../index.html)
[](https://github.com/ml-explore/mlx "Source repository")
* [.rst](../_sources/python/ops.rst "Download source file")
* .pdf
Operations
==========
Operations[#](#operations "Link to this heading")
==================================================
| | |
| --- | --- |
| [`abs`](_autosummary/mlx.core.abs.html#mlx.core.abs "mlx.core.abs")
(a, /, \*\[, stream\]) | Element-wise absolute value. |
| [`add`](_autosummary/mlx.core.add.html#mlx.core.add "mlx.core.add")
(a, b\[, stream\]) | Element-wise addition. |
| [`addmm`](_autosummary/mlx.core.addmm.html#mlx.core.addmm "mlx.core.addmm")
(c, a, b, /\[, alpha, beta, stream\]) | Matrix multiplication with addition and optional scaling. |
| [`all`](_autosummary/mlx.core.all.html#mlx.core.all "mlx.core.all")
(a, /\[, axis, keepdims, stream\]) | An and reduction over the given axes. |
| [`allclose`](_autosummary/mlx.core.allclose.html#mlx.core.allclose "mlx.core.allclose")
(a, b, /\[, rtol, atol, equal\_nan, ...\]) | Approximate comparison of two arrays. |
| [`any`](_autosummary/mlx.core.any.html#mlx.core.any "mlx.core.any")
(a, /\[, axis, keepdims, stream\]) | An or reduction over the given axes. |
| [`arange`](_autosummary/mlx.core.arange.html#mlx.core.arange "mlx.core.arange")
(-> array) | Overloaded function. |
| [`arccos`](_autosummary/mlx.core.arccos.html#mlx.core.arccos "mlx.core.arccos")
(a, /, \*\[, stream\]) | Element-wise inverse cosine. |
| [`arccosh`](_autosummary/mlx.core.arccosh.html#mlx.core.arccosh "mlx.core.arccosh")
(a, /, \*\[, stream\]) | Element-wise inverse hyperbolic cosine. |
| [`arcsin`](_autosummary/mlx.core.arcsin.html#mlx.core.arcsin "mlx.core.arcsin")
(a, /, \*\[, stream\]) | Element-wise inverse sine. |
| [`arcsinh`](_autosummary/mlx.core.arcsinh.html#mlx.core.arcsinh "mlx.core.arcsinh")
(a, /, \*\[, stream\]) | Element-wise inverse hyperbolic sine. |
| [`arctan`](_autosummary/mlx.core.arctan.html#mlx.core.arctan "mlx.core.arctan")
(a, /, \*\[, stream\]) | Element-wise inverse tangent. |
| [`arctan2`](_autosummary/mlx.core.arctan2.html#mlx.core.arctan2 "mlx.core.arctan2")
(a, b, /, \*\[, stream\]) | Element-wise inverse tangent of the ratio of two arrays. |
| [`arctanh`](_autosummary/mlx.core.arctanh.html#mlx.core.arctanh "mlx.core.arctanh")
(a, /, \*\[, stream\]) | Element-wise inverse hyperbolic tangent. |
| [`argmax`](_autosummary/mlx.core.argmax.html#mlx.core.argmax "mlx.core.argmax")
(a, /\[, axis, keepdims, stream\]) | Indices of the maximum values along the axis. |
| [`argmin`](_autosummary/mlx.core.argmin.html#mlx.core.argmin "mlx.core.argmin")
(a, /\[, axis, keepdims, stream\]) | Indices of the minimum values along the axis. |
| [`argpartition`](_autosummary/mlx.core.argpartition.html#mlx.core.argpartition "mlx.core.argpartition")
(a, /, kth\[, axis, stream\]) | Returns the indices that partition the array. |
| [`argsort`](_autosummary/mlx.core.argsort.html#mlx.core.argsort "mlx.core.argsort")
(a, /\[, axis, stream\]) | Returns the indices that sort the array. |
| [`array_equal`](_autosummary/mlx.core.array_equal.html#mlx.core.array_equal "mlx.core.array_equal")
(a, b\[, equal\_nan, stream\]) | Array equality check. |
| [`as_strided`](_autosummary/mlx.core.as_strided.html#mlx.core.as_strided "mlx.core.as_strided")
(a, /\[, shape, strides, offset, ...\]) | Create a view into the array with the given shape and strides. |
| [`atleast_1d`](_autosummary/mlx.core.atleast_1d.html#mlx.core.atleast_1d "mlx.core.atleast_1d")
(\*arys\[, stream\]) | Convert all arrays to have at least one dimension. |
| [`atleast_2d`](_autosummary/mlx.core.atleast_2d.html#mlx.core.atleast_2d "mlx.core.atleast_2d")
(\*arys\[, stream\]) | Convert all arrays to have at least two dimensions. |
| [`atleast_3d`](_autosummary/mlx.core.atleast_3d.html#mlx.core.atleast_3d "mlx.core.atleast_3d")
(\*arys\[, stream\]) | Convert all arrays to have at least three dimensions. |
| [`bitwise_and`](_autosummary/mlx.core.bitwise_and.html#mlx.core.bitwise_and "mlx.core.bitwise_and")
(a, b\[, stream\]) | Element-wise bitwise and. |
| [`bitwise_invert`](_autosummary/mlx.core.bitwise_invert.html#mlx.core.bitwise_invert "mlx.core.bitwise_invert")
(a\[, stream\]) | Element-wise bitwise inverse. |
| [`bitwise_or`](_autosummary/mlx.core.bitwise_or.html#mlx.core.bitwise_or "mlx.core.bitwise_or")
(a, b\[, stream\]) | Element-wise bitwise or. |
| [`bitwise_xor`](_autosummary/mlx.core.bitwise_xor.html#mlx.core.bitwise_xor "mlx.core.bitwise_xor")
(a, b\[, stream\]) | Element-wise bitwise xor. |
| [`block_masked_mm`](_autosummary/mlx.core.block_masked_mm.html#mlx.core.block_masked_mm "mlx.core.block_masked_mm")
(a, b, /\[, block\_size, ...\]) | Matrix multiplication with block masking. |
| [`broadcast_to`](_autosummary/mlx.core.broadcast_to.html#mlx.core.broadcast_to "mlx.core.broadcast_to")
(a, /, shape, \*\[, stream\]) | Broadcast an array to the given shape. |
| [`ceil`](_autosummary/mlx.core.ceil.html#mlx.core.ceil "mlx.core.ceil")
(a, /, \*\[, stream\]) | Element-wise ceil. |
| [`clip`](_autosummary/mlx.core.clip.html#mlx.core.clip "mlx.core.clip")
(a, /, a\_min, a\_max, \*\[, stream\]) | Clip the values of the array between the given minimum and maximum. |
| [`concatenate`](_autosummary/mlx.core.concatenate.html#mlx.core.concatenate "mlx.core.concatenate")
(arrays\[, axis, stream\]) | Concatenate the arrays along the given axis. |
| [`conj`](_autosummary/mlx.core.conj.html#mlx.core.conj "mlx.core.conj")
(a, \*\[, stream\]) | Return the elementwise complex conjugate of the input. |
| [`conjugate`](_autosummary/mlx.core.conjugate.html#mlx.core.conjugate "mlx.core.conjugate")
(a, \*\[, stream\]) | Return the elementwise complex conjugate of the input. |
| [`convolve`](_autosummary/mlx.core.convolve.html#mlx.core.convolve "mlx.core.convolve")
(a, v, /\[, mode, stream\]) | The discrete convolution of 1D arrays. |
| [`conv1d`](_autosummary/mlx.core.conv1d.html#mlx.core.conv1d "mlx.core.conv1d")
(input, weight, /\[, stride, padding, ...\]) | 1D convolution over an input with several channels |
| [`conv2d`](_autosummary/mlx.core.conv2d.html#mlx.core.conv2d "mlx.core.conv2d")
(input, weight, /\[, stride, padding, ...\]) | 2D convolution over an input with several channels |
| [`conv3d`](_autosummary/mlx.core.conv3d.html#mlx.core.conv3d "mlx.core.conv3d")
(input, weight, /\[, stride, padding, ...\]) | 3D convolution over an input with several channels |
| [`conv_transpose1d`](_autosummary/mlx.core.conv_transpose1d.html#mlx.core.conv_transpose1d "mlx.core.conv_transpose1d")
(input, weight, /\[, stride, ...\]) | 1D transposed convolution over an input with several channels |
| [`conv_transpose2d`](_autosummary/mlx.core.conv_transpose2d.html#mlx.core.conv_transpose2d "mlx.core.conv_transpose2d")
(input, weight, /\[, stride, ...\]) | 2D transposed convolution over an input with several channels |
| [`conv_transpose3d`](_autosummary/mlx.core.conv_transpose3d.html#mlx.core.conv_transpose3d "mlx.core.conv_transpose3d")
(input, weight, /\[, stride, ...\]) | 3D transposed convolution over an input with several channels |
| [`conv_general`](_autosummary/mlx.core.conv_general.html#mlx.core.conv_general "mlx.core.conv_general")
(input, weight, /\[, stride, ...\]) | General convolution over an input with several channels |
| [`cos`](_autosummary/mlx.core.cos.html#mlx.core.cos "mlx.core.cos")
(a, /, \*\[, stream\]) | Element-wise cosine. |
| [`cosh`](_autosummary/mlx.core.cosh.html#mlx.core.cosh "mlx.core.cosh")
(a, /, \*\[, stream\]) | Element-wise hyperbolic cosine. |
| [`cummax`](_autosummary/mlx.core.cummax.html#mlx.core.cummax "mlx.core.cummax")
(a, /\[, axis, reverse, inclusive, stream\]) | Return the cumulative maximum of the elements along the given axis. |
| [`cummin`](_autosummary/mlx.core.cummin.html#mlx.core.cummin "mlx.core.cummin")
(a, /\[, axis, reverse, inclusive, stream\]) | Return the cumulative minimum of the elements along the given axis. |
| [`cumprod`](_autosummary/mlx.core.cumprod.html#mlx.core.cumprod "mlx.core.cumprod")
(a, /\[, axis, reverse, inclusive, stream\]) | Return the cumulative product of the elements along the given axis. |
| [`cumsum`](_autosummary/mlx.core.cumsum.html#mlx.core.cumsum "mlx.core.cumsum")
(a, /\[, axis, reverse, inclusive, stream\]) | Return the cumulative sum of the elements along the given axis. |
| [`degrees`](_autosummary/mlx.core.degrees.html#mlx.core.degrees "mlx.core.degrees")
(a, /, \*\[, stream\]) | Convert angles from radians to degrees. |
| [`dequantize`](_autosummary/mlx.core.dequantize.html#mlx.core.dequantize "mlx.core.dequantize")
(w, /, scales, biases\[, ...\]) | Dequantize the matrix `w` using the provided `scales` and `biases` and the `group_size` and `bits` configuration. |
| [`diag`](_autosummary/mlx.core.diag.html#mlx.core.diag "mlx.core.diag")
(a, /\[, k, stream\]) | Extract a diagonal or construct a diagonal matrix. |
| [`diagonal`](_autosummary/mlx.core.diagonal.html#mlx.core.diagonal "mlx.core.diagonal")
(a\[, offset, axis1, axis2, stream\]) | Return specified diagonals. |
| [`divide`](_autosummary/mlx.core.divide.html#mlx.core.divide "mlx.core.divide")
(a, b\[, stream\]) | Element-wise division. |
| [`divmod`](_autosummary/mlx.core.divmod.html#mlx.core.divmod "mlx.core.divmod")
(a, b\[, stream\]) | Element-wise quotient and remainder. |
| [`einsum`](_autosummary/mlx.core.einsum.html#mlx.core.einsum "mlx.core.einsum")
(subscripts, \*operands\[, stream\]) | Perform the Einstein summation convention on the operands. |
| [`einsum_path`](_autosummary/mlx.core.einsum_path.html#mlx.core.einsum_path "mlx.core.einsum_path")
(subscripts, \*operands) | Compute the contraction order for the given Einstein summation. |
| [`equal`](_autosummary/mlx.core.equal.html#mlx.core.equal "mlx.core.equal")
(a, b\[, stream\]) | Element-wise equality. |
| [`erf`](_autosummary/mlx.core.erf.html#mlx.core.erf "mlx.core.erf")
(a, /, \*\[, stream\]) | Element-wise error function. |
| [`erfinv`](_autosummary/mlx.core.erfinv.html#mlx.core.erfinv "mlx.core.erfinv")
(a, /, \*\[, stream\]) | Element-wise inverse of [`erf()`](_autosummary/mlx.core.erf.html#mlx.core.erf "mlx.core.erf")
. |
| [`exp`](_autosummary/mlx.core.exp.html#mlx.core.exp "mlx.core.exp")
(a, /, \*\[, stream\]) | Element-wise exponential. |
| [`expm1`](_autosummary/mlx.core.expm1.html#mlx.core.expm1 "mlx.core.expm1")
(a, /, \*\[, stream\]) | Element-wise exponential minus 1. |
| [`expand_dims`](_autosummary/mlx.core.expand_dims.html#mlx.core.expand_dims "mlx.core.expand_dims")
(a, /, axis, \*\[, stream\]) | Add a size one dimension at the given axis. |
| [`eye`](_autosummary/mlx.core.eye.html#mlx.core.eye "mlx.core.eye")
(n\[, m, k, dtype, stream\]) | Create an identity matrix or a general diagonal matrix. |
| [`flatten`](_autosummary/mlx.core.flatten.html#mlx.core.flatten "mlx.core.flatten")
(a, /\[, start\_axis, end\_axis, stream\]) | Flatten an array. |
| [`floor`](_autosummary/mlx.core.floor.html#mlx.core.floor "mlx.core.floor")
(a, /, \*\[, stream\]) | Element-wise floor. |
| [`floor_divide`](_autosummary/mlx.core.floor_divide.html#mlx.core.floor_divide "mlx.core.floor_divide")
(a, b\[, stream\]) | Element-wise integer division. |
| [`full`](_autosummary/mlx.core.full.html#mlx.core.full "mlx.core.full")
(shape, vals\[, dtype, stream\]) | Construct an array with the given value. |
| [`gather_mm`](_autosummary/mlx.core.gather_mm.html#mlx.core.gather_mm "mlx.core.gather_mm")
(a, b, /, lhs\_indices, rhs\_indices, \*) | Matrix multiplication with matrix-level gather. |
| [`gather_qmm`](_autosummary/mlx.core.gather_qmm.html#mlx.core.gather_qmm "mlx.core.gather_qmm")
(x, w, /, scales, biases\[, ...\]) | Perform quantized matrix multiplication with matrix-level gather. |
| [`greater`](_autosummary/mlx.core.greater.html#mlx.core.greater "mlx.core.greater")
(a, b\[, stream\]) | Element-wise greater than. |
| [`greater_equal`](_autosummary/mlx.core.greater_equal.html#mlx.core.greater_equal "mlx.core.greater_equal")
(a, b\[, stream\]) | Element-wise greater or equal. |
| [`hadamard_transform`](_autosummary/mlx.core.hadamard_transform.html#mlx.core.hadamard_transform "mlx.core.hadamard_transform")
(a\[, scale, stream\]) | Perform the Walsh-Hadamard transform along the final axis. |
| [`identity`](_autosummary/mlx.core.identity.html#mlx.core.identity "mlx.core.identity")
(n\[, dtype, stream\]) | Create a square identity matrix. |
| [`imag`](_autosummary/mlx.core.imag.html#mlx.core.imag "mlx.core.imag")
(a, /, \*\[, stream\]) | Returns the imaginary part of a complex array. |
| [`inner`](_autosummary/mlx.core.inner.html#mlx.core.inner "mlx.core.inner")
(a, b, /, \*\[, stream\]) | Ordinary inner product of vectors for 1-D arrays, in higher dimensions a sum product over the last axes. |
| [`isfinite`](_autosummary/mlx.core.isfinite.html#mlx.core.isfinite "mlx.core.isfinite")
(a\[, stream\]) | Return a boolean array indicating which elements are finite. |
| [`isclose`](_autosummary/mlx.core.isclose.html#mlx.core.isclose "mlx.core.isclose")
(a, b, /\[, rtol, atol, equal\_nan, stream\]) | Returns a boolean array where two arrays are element-wise equal within a tolerance. |
| [`isinf`](_autosummary/mlx.core.isinf.html#mlx.core.isinf "mlx.core.isinf")
(a\[, stream\]) | Return a boolean array indicating which elements are +/- inifnity. |
| [`isnan`](_autosummary/mlx.core.isnan.html#mlx.core.isnan "mlx.core.isnan")
(a\[, stream\]) | Return a boolean array indicating which elements are NaN. |
| [`isneginf`](_autosummary/mlx.core.isneginf.html#mlx.core.isneginf "mlx.core.isneginf")
(a\[, stream\]) | Return a boolean array indicating which elements are negative infinity. |
| [`isposinf`](_autosummary/mlx.core.isposinf.html#mlx.core.isposinf "mlx.core.isposinf")
(a\[, stream\]) | Return a boolean array indicating which elements are positive infinity. |
| [`issubdtype`](_autosummary/mlx.core.issubdtype.html#mlx.core.issubdtype "mlx.core.issubdtype")
(arg1, arg2) | Check if a [`Dtype`](_autosummary/mlx.core.Dtype.html#mlx.core.Dtype "mlx.core.Dtype")
or [`DtypeCategory`](_autosummary/mlx.core.DtypeCategory.html#mlx.core.DtypeCategory "mlx.core.DtypeCategory")
is a subtype of another. |
| [`kron`](_autosummary/mlx.core.kron.html#mlx.core.kron "mlx.core.kron")
(a, b, \*\[, stream\]) | Compute the Kronecker product of two arrays `a` and `b`. |
| [`left_shift`](_autosummary/mlx.core.left_shift.html#mlx.core.left_shift "mlx.core.left_shift")
(a, b\[, stream\]) | Element-wise left shift. |
| [`less`](_autosummary/mlx.core.less.html#mlx.core.less "mlx.core.less")
(a, b\[, stream\]) | Element-wise less than. |
| [`less_equal`](_autosummary/mlx.core.less_equal.html#mlx.core.less_equal "mlx.core.less_equal")
(a, b\[, stream\]) | Element-wise less than or equal. |
| [`linspace`](_autosummary/mlx.core.linspace.html#mlx.core.linspace "mlx.core.linspace")
(start, stop\[, num, dtype, stream\]) | Generate `num` evenly spaced numbers over interval `[start, stop]`. |
| [`load`](_autosummary/mlx.core.load.html#mlx.core.load "mlx.core.load")
(file, /\[, format, return\_metadata, stream\]) | Load array(s) from a binary file. |
| [`log`](_autosummary/mlx.core.log.html#mlx.core.log "mlx.core.log")
(a, /, \*\[, stream\]) | Element-wise natural logarithm. |
| [`log2`](_autosummary/mlx.core.log2.html#mlx.core.log2 "mlx.core.log2")
(a, /, \*\[, stream\]) | Element-wise base-2 logarithm. |
| [`log10`](_autosummary/mlx.core.log10.html#mlx.core.log10 "mlx.core.log10")
(a, /, \*\[, stream\]) | Element-wise base-10 logarithm. |
| [`log1p`](_autosummary/mlx.core.log1p.html#mlx.core.log1p "mlx.core.log1p")
(a, /, \*\[, stream\]) | Element-wise natural log of one plus the array. |
| [`logaddexp`](_autosummary/mlx.core.logaddexp.html#mlx.core.logaddexp "mlx.core.logaddexp")
(a, b, /, \*\[, stream\]) | Element-wise log-add-exp. |
| [`logical_not`](_autosummary/mlx.core.logical_not.html#mlx.core.logical_not "mlx.core.logical_not")
(a, /, \*\[, stream\]) | Element-wise logical not. |
| [`logical_and`](_autosummary/mlx.core.logical_and.html#mlx.core.logical_and "mlx.core.logical_and")
(a, b, /, \*\[, stream\]) | Element-wise logical and. |
| [`logical_or`](_autosummary/mlx.core.logical_or.html#mlx.core.logical_or "mlx.core.logical_or")
(a, b, /, \*\[, stream\]) | Element-wise logical or. |
| [`logsumexp`](_autosummary/mlx.core.logsumexp.html#mlx.core.logsumexp "mlx.core.logsumexp")
(a, /\[, axis, keepdims, stream\]) | A log-sum-exp reduction over the given axes. |
| [`matmul`](_autosummary/mlx.core.matmul.html#mlx.core.matmul "mlx.core.matmul")
(a, b, /, \*\[, stream\]) | Matrix multiplication. |
| [`max`](_autosummary/mlx.core.max.html#mlx.core.max "mlx.core.max")
(a, /\[, axis, keepdims, stream\]) | A max reduction over the given axes. |
| [`maximum`](_autosummary/mlx.core.maximum.html#mlx.core.maximum "mlx.core.maximum")
(a, b, /, \*\[, stream\]) | Element-wise maximum. |
| [`mean`](_autosummary/mlx.core.mean.html#mlx.core.mean "mlx.core.mean")
(a, /\[, axis, keepdims, stream\]) | Compute the mean(s) over the given axes. |
| [`meshgrid`](_autosummary/mlx.core.meshgrid.html#mlx.core.meshgrid "mlx.core.meshgrid")
(\*arrays\[, sparse, indexing, stream\]) | Generate multidimensional coordinate grids from 1-D coordinate arrays |
| [`min`](_autosummary/mlx.core.min.html#mlx.core.min "mlx.core.min")
(a, /\[, axis, keepdims, stream\]) | A min reduction over the given axes. |
| [`minimum`](_autosummary/mlx.core.minimum.html#mlx.core.minimum "mlx.core.minimum")
(a, b, /, \*\[, stream\]) | Element-wise minimum. |
| [`moveaxis`](_autosummary/mlx.core.moveaxis.html#mlx.core.moveaxis "mlx.core.moveaxis")
(a, /, source, destination, \*\[, stream\]) | Move an axis to a new position. |
| [`multiply`](_autosummary/mlx.core.multiply.html#mlx.core.multiply "mlx.core.multiply")
(a, b\[, stream\]) | Element-wise multiplication. |
| [`nan_to_num`](_autosummary/mlx.core.nan_to_num.html#mlx.core.nan_to_num "mlx.core.nan_to_num")
(a\[, nan, posinf, neginf, stream\]) | Replace NaN and Inf values with finite numbers. |
| [`negative`](_autosummary/mlx.core.negative.html#mlx.core.negative "mlx.core.negative")
(a, /, \*\[, stream\]) | Element-wise negation. |
| [`not_equal`](_autosummary/mlx.core.not_equal.html#mlx.core.not_equal "mlx.core.not_equal")
(a, b\[, stream\]) | Element-wise not equal. |
| [`ones`](_autosummary/mlx.core.ones.html#mlx.core.ones "mlx.core.ones")
(shape\[, dtype, stream\]) | Construct an array of ones. |
| [`ones_like`](_autosummary/mlx.core.ones_like.html#mlx.core.ones_like "mlx.core.ones_like")
(a, /, \*\[, stream\]) | An array of ones like the input. |
| [`outer`](_autosummary/mlx.core.outer.html#mlx.core.outer "mlx.core.outer")
(a, b, /, \*\[, stream\]) | Compute the outer product of two 1-D arrays, if the array's passed are not 1-D a flatten op will be run beforehand. |
| [`partition`](_autosummary/mlx.core.partition.html#mlx.core.partition "mlx.core.partition")
(a, /, kth\[, axis, stream\]) | Returns a partitioned copy of the array such that the smaller `kth` elements are first. |
| [`pad`](_autosummary/mlx.core.pad.html#mlx.core.pad "mlx.core.pad")
(a, pad\_width\[, mode, constant\_values, ...\]) | Pad an array with a constant value |
| [`power`](_autosummary/mlx.core.power.html#mlx.core.power "mlx.core.power")
(a, b, /, \*\[, stream\]) | Element-wise power operation. |
| [`prod`](_autosummary/mlx.core.prod.html#mlx.core.prod "mlx.core.prod")
(a, /\[, axis, keepdims, stream\]) | An product reduction over the given axes. |
| [`put_along_axis`](_autosummary/mlx.core.put_along_axis.html#mlx.core.put_along_axis "mlx.core.put_along_axis")
(a, /, indices, values\[, ...\]) | Put values along an axis at the specified indices. |
| [`quantize`](_autosummary/mlx.core.quantize.html#mlx.core.quantize "mlx.core.quantize")
(w, /\[, group\_size, bits, stream\]) | Quantize the matrix `w` using `bits` bits per element. |
| [`quantized_matmul`](_autosummary/mlx.core.quantized_matmul.html#mlx.core.quantized_matmul "mlx.core.quantized_matmul")
(x, w, /, scales, biases\[, ...\]) | Perform the matrix multiplication with the quantized matrix `w`. |
| [`radians`](_autosummary/mlx.core.radians.html#mlx.core.radians "mlx.core.radians")
(a, /, \*\[, stream\]) | Convert angles from degrees to radians. |
| [`real`](_autosummary/mlx.core.real.html#mlx.core.real "mlx.core.real")
(a, /, \*\[, stream\]) | Returns the real part of a complex array. |
| [`reciprocal`](_autosummary/mlx.core.reciprocal.html#mlx.core.reciprocal "mlx.core.reciprocal")
(a, /, \*\[, stream\]) | Element-wise reciprocal. |
| [`remainder`](_autosummary/mlx.core.remainder.html#mlx.core.remainder "mlx.core.remainder")
(a, b\[, stream\]) | Element-wise remainder of division. |
| [`repeat`](_autosummary/mlx.core.repeat.html#mlx.core.repeat "mlx.core.repeat")
(array, repeats\[, axis, stream\]) | Repeat an array along a specified axis. |
| [`reshape`](_autosummary/mlx.core.reshape.html#mlx.core.reshape "mlx.core.reshape")
(a, /, shape, \*\[, stream\]) | Reshape an array while preserving the size. |
| [`right_shift`](_autosummary/mlx.core.right_shift.html#mlx.core.right_shift "mlx.core.right_shift")
(a, b\[, stream\]) | Element-wise right shift. |
| [`roll`](_autosummary/mlx.core.roll.html#mlx.core.roll "mlx.core.roll")
(a, shift\[, axis, stream\]) | Roll array elements along a given axis. |
| [`round`](_autosummary/mlx.core.round.html#mlx.core.round "mlx.core.round")
(a, /\[, decimals, stream\]) | Round to the given number of decimals. |
| [`rsqrt`](_autosummary/mlx.core.rsqrt.html#mlx.core.rsqrt "mlx.core.rsqrt")
(a, /, \*\[, stream\]) | Element-wise reciprocal and square root. |
| [`save`](_autosummary/mlx.core.save.html#mlx.core.save "mlx.core.save")
(file, arr) | Save the array to a binary file in `.npy` format. |
| [`savez`](_autosummary/mlx.core.savez.html#mlx.core.savez "mlx.core.savez")
(file, \*args, \*\*kwargs) | Save several arrays to a binary file in uncompressed `.npz` format. |
| [`savez_compressed`](_autosummary/mlx.core.savez_compressed.html#mlx.core.savez_compressed "mlx.core.savez_compressed")
(file, \*args, \*\*kwargs) | Save several arrays to a binary file in compressed `.npz` format. |
| [`save_gguf`](_autosummary/mlx.core.save_gguf.html#mlx.core.save_gguf "mlx.core.save_gguf")
(file, arrays, metadata) | Save array(s) to a binary file in `.gguf` format. |
| [`save_safetensors`](_autosummary/mlx.core.save_safetensors.html#mlx.core.save_safetensors "mlx.core.save_safetensors")
(file, arrays\[, metadata\]) | Save array(s) to a binary file in `.safetensors` format. |
| [`sigmoid`](_autosummary/mlx.core.sigmoid.html#mlx.core.sigmoid "mlx.core.sigmoid")
(a, /, \*\[, stream\]) | Element-wise logistic sigmoid. |
| [`sign`](_autosummary/mlx.core.sign.html#mlx.core.sign "mlx.core.sign")
(a, /, \*\[, stream\]) | Element-wise sign. |
| [`sin`](_autosummary/mlx.core.sin.html#mlx.core.sin "mlx.core.sin")
(a, /, \*\[, stream\]) | Element-wise sine. |
| [`sinh`](_autosummary/mlx.core.sinh.html#mlx.core.sinh "mlx.core.sinh")
(a, /, \*\[, stream\]) | Element-wise hyperbolic sine. |
| [`slice`](_autosummary/mlx.core.slice.html#mlx.core.slice "mlx.core.slice")
(a, start\_indices, axes, slice\_size, \*) | Extract a sub-array from the input array. |
| [`slice_update`](_autosummary/mlx.core.slice_update.html#mlx.core.slice_update "mlx.core.slice_update")
(a, update, start\_indices, axes, \*) | Update a sub-array of the input array. |
| [`softmax`](_autosummary/mlx.core.softmax.html#mlx.core.softmax "mlx.core.softmax")
(a, /\[, axis, stream\]) | Perform the softmax along the given axis. |
| [`sort`](_autosummary/mlx.core.sort.html#mlx.core.sort "mlx.core.sort")
(a, /\[, axis, stream\]) | Returns a sorted copy of the array. |
| [`split`](_autosummary/mlx.core.split.html#mlx.core.split "mlx.core.split")
(a, /, indices\_or\_sections\[, axis, stream\]) | Split an array along a given axis. |
| [`sqrt`](_autosummary/mlx.core.sqrt.html#mlx.core.sqrt "mlx.core.sqrt")
(a, /, \*\[, stream\]) | Element-wise square root. |
| [`square`](_autosummary/mlx.core.square.html#mlx.core.square "mlx.core.square")
(a, /, \*\[, stream\]) | Element-wise square. |
| [`squeeze`](_autosummary/mlx.core.squeeze.html#mlx.core.squeeze "mlx.core.squeeze")
(a, /\[, axis, stream\]) | Remove length one axes from an array. |
| [`stack`](_autosummary/mlx.core.stack.html#mlx.core.stack "mlx.core.stack")
(arrays\[, axis, stream\]) | Stacks the arrays along a new axis. |
| [`std`](_autosummary/mlx.core.std.html#mlx.core.std "mlx.core.std")
(a, /\[, axis, keepdims, ddof, stream\]) | Compute the standard deviation(s) over the given axes. |
| [`stop_gradient`](_autosummary/mlx.core.stop_gradient.html#mlx.core.stop_gradient "mlx.core.stop_gradient")
(a, /, \*\[, stream\]) | Stop gradients from being computed. |
| [`subtract`](_autosummary/mlx.core.subtract.html#mlx.core.subtract "mlx.core.subtract")
(a, b\[, stream\]) | Element-wise subtraction. |
| [`sum`](_autosummary/mlx.core.sum.html#mlx.core.sum "mlx.core.sum")
(a, /\[, axis, keepdims, stream\]) | Sum reduce the array over the given axes. |
| [`swapaxes`](_autosummary/mlx.core.swapaxes.html#mlx.core.swapaxes "mlx.core.swapaxes")
(a, /, axis1, axis2, \*\[, stream\]) | Swap two axes of an array. |
| [`take`](_autosummary/mlx.core.take.html#mlx.core.take "mlx.core.take")
(a, /, indices\[, axis, stream\]) | Take elements along an axis. |
| [`take_along_axis`](_autosummary/mlx.core.take_along_axis.html#mlx.core.take_along_axis "mlx.core.take_along_axis")
(a, /, indices\[, axis, stream\]) | Take values along an axis at the specified indices. |
| [`tan`](_autosummary/mlx.core.tan.html#mlx.core.tan "mlx.core.tan")
(a, /, \*\[, stream\]) | Element-wise tangent. |
| [`tanh`](_autosummary/mlx.core.tanh.html#mlx.core.tanh "mlx.core.tanh")
(a, /, \*\[, stream\]) | Element-wise hyperbolic tangent. |
| [`tensordot`](_autosummary/mlx.core.tensordot.html#mlx.core.tensordot "mlx.core.tensordot")
(a, b, /\[, axes, stream\]) | Compute the tensor dot product along the specified axes. |
| [`tile`](_autosummary/mlx.core.tile.html#mlx.core.tile "mlx.core.tile")
(a, reps, /, \*\[, stream\]) | Construct an array by repeating `a` the number of times given by `reps`. |
| [`topk`](_autosummary/mlx.core.topk.html#mlx.core.topk "mlx.core.topk")
(a, /, k\[, axis, stream\]) | Returns the `k` largest elements from the input along a given axis. |
| [`trace`](_autosummary/mlx.core.trace.html#mlx.core.trace "mlx.core.trace")
(a, /\[, offset, axis1, axis2, dtype, ...\]) | Return the sum along a specified diagonal in the given array. |
| [`transpose`](_autosummary/mlx.core.transpose.html#mlx.core.transpose "mlx.core.transpose")
(a, /\[, axes, stream\]) | Transpose the dimensions of the array. |
| [`tri`](_autosummary/mlx.core.tri.html#mlx.core.tri "mlx.core.tri")
(n, m, k\[, dtype, stream\]) | An array with ones at and below the given diagonal and zeros elsewhere. |
| [`tril`](_autosummary/mlx.core.tril.html#mlx.core.tril "mlx.core.tril")
(x, k, \*\[, stream\]) | Zeros the array above the given diagonal. |
| [`triu`](_autosummary/mlx.core.triu.html#mlx.core.triu "mlx.core.triu")
(x, k, \*\[, stream\]) | Zeros the array below the given diagonal. |
| [`unflatten`](_autosummary/mlx.core.unflatten.html#mlx.core.unflatten "mlx.core.unflatten")
(a, /, axis, shape, \*\[, stream\]) | Unflatten an axis of an array to a shape. |
| [`var`](_autosummary/mlx.core.var.html#mlx.core.var "mlx.core.var")
(a, /\[, axis, keepdims, ddof, stream\]) | Compute the variance(s) over the given axes. |
| [`view`](_autosummary/mlx.core.view.html#mlx.core.view "mlx.core.view")
(a, dtype\[, stream\]) | View the array as a different type. |
| [`where`](_autosummary/mlx.core.where.html#mlx.core.where "mlx.core.where")
(condition, x, y, /, \*\[, stream\]) | Select from `x` or `y` according to `condition`. |
| [`zeros`](_autosummary/mlx.core.zeros.html#mlx.core.zeros "mlx.core.zeros")
(shape\[, dtype, stream\]) | Construct an array of zeros. |
| [`zeros_like`](_autosummary/mlx.core.zeros_like.html#mlx.core.zeros_like "mlx.core.zeros_like")
(a, /, \*\[, stream\]) | An array of zeros like the input. |
---
# Fast — MLX 0.24.0 documentation
[Skip to main content](#main-content)
Back to top
Ctrl+K
[](../index.html)
[](https://github.com/ml-explore/mlx "Source repository")
* [.rst](../_sources/python/fast.rst "Download source file")
* .pdf
Fast
====
Fast[#](#fast "Link to this heading")
======================================
| | |
| --- | --- |
| [`rms_norm`](_autosummary/mlx.core.fast.rms_norm.html#mlx.core.fast.rms_norm "mlx.core.fast.rms_norm")
(x, weight, eps, \*\[, stream\]) | Root Mean Square normalization (RMS norm). |
| [`layer_norm`](_autosummary/mlx.core.fast.layer_norm.html#mlx.core.fast.layer_norm "mlx.core.fast.layer_norm")
(x, weight, bias, eps, \*\[, stream\]) | Layer normalization. |
| [`rope`](_autosummary/mlx.core.fast.rope.html#mlx.core.fast.rope "mlx.core.fast.rope")
(a, dims, \*, traditional, base, scale, ...) | Apply rotary positional encoding to the input. |
| [`scaled_dot_product_attention`](_autosummary/mlx.core.fast.scaled_dot_product_attention.html#mlx.core.fast.scaled_dot_product_attention "mlx.core.fast.scaled_dot_product_attention")
(q, k, v, \*, scale) | A fast implementation of multi-head attention: `O = softmax(Q @ K.T, dim=-1) @ V`. |
| [`metal_kernel`](_autosummary/mlx.core.fast.metal_kernel.html#mlx.core.fast.metal_kernel "mlx.core.fast.metal_kernel")
(name, input\_names, ...\[, ...\]) | A jit-compiled custom Metal kernel defined from a source string. |
---
# Transforms — MLX 0.24.0 documentation
[Skip to main content](#main-content)
Back to top
Ctrl+K
[](../index.html)
[](https://github.com/ml-explore/mlx "Source repository")
* [.rst](../_sources/python/transforms.rst "Download source file")
* .pdf
Transforms
==========
Transforms[#](#transforms "Link to this heading")
==================================================
| | |
| --- | --- |
| [`eval`](_autosummary/mlx.core.eval.html#mlx.core.eval "mlx.core.eval")
(\*args) | Evaluate an [`array`](_autosummary/mlx.core.array.html#mlx.core.array "mlx.core.array")
or tree of [`array`](_autosummary/mlx.core.array.html#mlx.core.array "mlx.core.array")
. |
| [`compile`](_autosummary/mlx.core.compile.html#mlx.core.compile "mlx.core.compile")
(fun\[, inputs, outputs, shapeless\]) | Returns a compiled function which produces the same output as `fun`. |
| [`custom_function`](_autosummary/mlx.core.custom_function.html#mlx.core.custom_function "mlx.core.custom_function") | Set up a function for custom gradient and vmap definitions. |
| [`disable_compile`](_autosummary/mlx.core.disable_compile.html#mlx.core.disable_compile "mlx.core.disable_compile")
() | Globally disable compilation. |
| [`enable_compile`](_autosummary/mlx.core.enable_compile.html#mlx.core.enable_compile "mlx.core.enable_compile")
() | Globally enable compilation. |
| [`grad`](_autosummary/mlx.core.grad.html#mlx.core.grad "mlx.core.grad")
(fun\[, argnums, argnames\]) | Returns a function which computes the gradient of `fun`. |
| [`value_and_grad`](_autosummary/mlx.core.value_and_grad.html#mlx.core.value_and_grad "mlx.core.value_and_grad")
(fun\[, argnums, argnames\]) | Returns a function which computes the value and gradient of `fun`. |
| [`jvp`](_autosummary/mlx.core.jvp.html#mlx.core.jvp "mlx.core.jvp")
(fun, primals, tangents) | Compute the Jacobian-vector product. |
| [`vjp`](_autosummary/mlx.core.vjp.html#mlx.core.vjp "mlx.core.vjp")
(fun, primals, cotangents) | Compute the vector-Jacobian product. |
| [`vmap`](_autosummary/mlx.core.vmap.html#mlx.core.vmap "mlx.core.vmap")
(fun\[, in\_axes, out\_axes\]) | Returns a vectorized version of `fun`. |
---
# FFT — MLX 0.24.0 documentation
[Skip to main content](#main-content)
Back to top
Ctrl+K
[](../index.html)
[](https://github.com/ml-explore/mlx "Source repository")
* [.rst](../_sources/python/fft.rst "Download source file")
* .pdf
FFT
===
FFT[#](#fft "Link to this heading")
====================================
| | |
| --- | --- |
| [`fft`](_autosummary/mlx.core.fft.fft.html#mlx.core.fft.fft "mlx.core.fft.fft")
(a\[, n, axis, stream\]) | One dimensional discrete Fourier Transform. |
| [`ifft`](_autosummary/mlx.core.fft.ifft.html#mlx.core.fft.ifft "mlx.core.fft.ifft")
(a\[, n, axis, stream\]) | One dimensional inverse discrete Fourier Transform. |
| [`fft2`](_autosummary/mlx.core.fft.fft2.html#mlx.core.fft.fft2 "mlx.core.fft.fft2")
(a\[, s, axes, stream\]) | Two dimensional discrete Fourier Transform. |
| [`ifft2`](_autosummary/mlx.core.fft.ifft2.html#mlx.core.fft.ifft2 "mlx.core.fft.ifft2")
(a\[, s, axes, stream\]) | Two dimensional inverse discrete Fourier Transform. |
| [`fftn`](_autosummary/mlx.core.fft.fftn.html#mlx.core.fft.fftn "mlx.core.fft.fftn")
(a\[, s, axes, stream\]) | n-dimensional discrete Fourier Transform. |
| [`ifftn`](_autosummary/mlx.core.fft.ifftn.html#mlx.core.fft.ifftn "mlx.core.fft.ifftn")
(a\[, s, axes, stream\]) | n-dimensional inverse discrete Fourier Transform. |
| [`rfft`](_autosummary/mlx.core.fft.rfft.html#mlx.core.fft.rfft "mlx.core.fft.rfft")
(a\[, n, axis, stream\]) | One dimensional discrete Fourier Transform on a real input. |
| [`irfft`](_autosummary/mlx.core.fft.irfft.html#mlx.core.fft.irfft "mlx.core.fft.irfft")
(a\[, n, axis, stream\]) | The inverse of [`rfft()`](_autosummary/mlx.core.fft.rfft.html#mlx.core.fft.rfft "mlx.core.fft.rfft")
. |
| [`rfft2`](_autosummary/mlx.core.fft.rfft2.html#mlx.core.fft.rfft2 "mlx.core.fft.rfft2")
(a\[, s, axes, stream\]) | Two dimensional real discrete Fourier Transform. |
| [`irfft2`](_autosummary/mlx.core.fft.irfft2.html#mlx.core.fft.irfft2 "mlx.core.fft.irfft2")
(a\[, s, axes, stream\]) | The inverse of [`rfft2()`](_autosummary/mlx.core.fft.rfft2.html#mlx.core.fft.rfft2 "mlx.core.fft.rfft2")
. |
| [`rfftn`](_autosummary/mlx.core.fft.rfftn.html#mlx.core.fft.rfftn "mlx.core.fft.rfftn")
(a\[, s, axes, stream\]) | n-dimensional real discrete Fourier Transform. |
| [`irfftn`](_autosummary/mlx.core.fft.irfftn.html#mlx.core.fft.irfftn "mlx.core.fft.irfftn")
(a\[, s, axes, stream\]) | The inverse of [`rfftn()`](_autosummary/mlx.core.fft.rfftn.html#mlx.core.fft.rfftn "mlx.core.fft.rfftn")
. |
---
# Linear Algebra — MLX 0.24.0 documentation
[Skip to main content](#main-content)
Back to top
Ctrl+K
[](../index.html)
[](https://github.com/ml-explore/mlx "Source repository")
* [.rst](../_sources/python/linalg.rst "Download source file")
* .pdf
Linear Algebra
==============
Linear Algebra[#](#linear-algebra "Link to this heading")
==========================================================
| | |
| --- | --- |
| [`inv`](_autosummary/mlx.core.linalg.inv.html#mlx.core.linalg.inv "mlx.core.linalg.inv")
(a, \*\[, stream\]) | Compute the inverse of a square matrix. |
| [`tri_inv`](_autosummary/mlx.core.linalg.tri_inv.html#mlx.core.linalg.tri_inv "mlx.core.linalg.tri_inv")
(a\[, upper, stream\]) | Compute the inverse of a triangular square matrix. |
| [`norm`](_autosummary/mlx.core.linalg.norm.html#mlx.core.linalg.norm "mlx.core.linalg.norm")
(a, /\[, ord, axis, keepdims, stream\]) | Matrix or vector norm. |
| [`cholesky`](_autosummary/mlx.core.linalg.cholesky.html#mlx.core.linalg.cholesky "mlx.core.linalg.cholesky")
(a\[, upper, stream\]) | Compute the Cholesky decomposition of a real symmetric positive semi-definite matrix. |
| [`cholesky_inv`](_autosummary/mlx.core.linalg.cholesky_inv.html#mlx.core.linalg.cholesky_inv "mlx.core.linalg.cholesky_inv")
(L\[, upper, stream\]) | Compute the inverse of a real symmetric positive semi-definite matrix using it's Cholesky decomposition. |
| [`cross`](_autosummary/mlx.core.linalg.cross.html#mlx.core.linalg.cross "mlx.core.linalg.cross")
(a, b\[, axis, stream\]) | Compute the cross product of two arrays along a specified axis. |
| [`qr`](_autosummary/mlx.core.linalg.qr.html#mlx.core.linalg.qr "mlx.core.linalg.qr")
(a, \*\[, stream\]) | The QR factorization of the input matrix. |
| [`svd`](_autosummary/mlx.core.linalg.svd.html#mlx.core.linalg.svd "mlx.core.linalg.svd")
(a\[, compute\_uv, stream\]) | The Singular Value Decomposition (SVD) of the input matrix. |
| [`eigvalsh`](_autosummary/mlx.core.linalg.eigvalsh.html#mlx.core.linalg.eigvalsh "mlx.core.linalg.eigvalsh")
(a\[, UPLO, stream\]) | Compute the eigenvalues of a complex Hermitian or real symmetric matrix. |
| [`eigh`](_autosummary/mlx.core.linalg.eigh.html#mlx.core.linalg.eigh "mlx.core.linalg.eigh")
(a\[, UPLO, stream\]) | Compute the eigenvalues and eigenvectors of a complex Hermitian or real symmetric matrix. |
| [`lu`](_autosummary/mlx.core.linalg.lu.html#mlx.core.linalg.lu "mlx.core.linalg.lu")
(a, \*\[, stream\]) | Compute the LU factorization of the given matrix `A`. |
| [`lu_factor`](_autosummary/mlx.core.linalg.lu_factor.html#mlx.core.linalg.lu_factor "mlx.core.linalg.lu_factor")
(a, \*\[, stream\]) | Computes a compact representation of the LU factorization. |
| [`solve`](_autosummary/mlx.core.linalg.solve.html#mlx.core.linalg.solve "mlx.core.linalg.solve")
(a, b, \*\[, stream\]) | Compute the solution to a system of linear equations `AX = B`. |
| [`solve_triangular`](_autosummary/mlx.core.linalg.solve_triangular.html#mlx.core.linalg.solve_triangular "mlx.core.linalg.solve_triangular")
(a, b, \*\[, upper, stream\]) | Computes the solution of a triangular system of linear equations `AX = B`. |
---
# Metal — MLX 0.24.0 documentation
[Skip to main content](#main-content)
Back to top
Ctrl+K
[](../index.html)
[](https://github.com/ml-explore/mlx "Source repository")
* [.rst](../_sources/python/metal.rst "Download source file")
* .pdf
Metal
=====
Metal[#](#metal "Link to this heading")
========================================
| | |
| --- | --- |
| [`is_available`](_autosummary/mlx.core.metal.is_available.html#mlx.core.metal.is_available "mlx.core.metal.is_available")
() | Check if the Metal back-end is available. |
| [`device_info`](_autosummary/mlx.core.metal.device_info.html#mlx.core.metal.device_info "mlx.core.metal.device_info")
() | Get information about the GPU device and system settings. |
| [`get_active_memory`](_autosummary/mlx.core.metal.get_active_memory.html#mlx.core.metal.get_active_memory "mlx.core.metal.get_active_memory")
() | Get the actively used memory in bytes. |
| [`get_peak_memory`](_autosummary/mlx.core.metal.get_peak_memory.html#mlx.core.metal.get_peak_memory "mlx.core.metal.get_peak_memory")
() | Get the peak amount of used memory in bytes. |
| [`reset_peak_memory`](_autosummary/mlx.core.metal.reset_peak_memory.html#mlx.core.metal.reset_peak_memory "mlx.core.metal.reset_peak_memory")
() | Reset the peak memory to zero. |
| [`get_cache_memory`](_autosummary/mlx.core.metal.get_cache_memory.html#mlx.core.metal.get_cache_memory "mlx.core.metal.get_cache_memory")
() | Get the cache size in bytes. |
| [`set_memory_limit`](_autosummary/mlx.core.metal.set_memory_limit.html#mlx.core.metal.set_memory_limit "mlx.core.metal.set_memory_limit")
(limit, \*\[, relaxed\]) | Set the memory limit. |
| [`set_cache_limit`](_autosummary/mlx.core.metal.set_cache_limit.html#mlx.core.metal.set_cache_limit "mlx.core.metal.set_cache_limit")
(limit) | Set the free cache limit. |
| [`set_wired_limit`](_autosummary/mlx.core.metal.set_wired_limit.html#mlx.core.metal.set_wired_limit "mlx.core.metal.set_wired_limit")
(limit) | Set the wired size limit. |
| [`clear_cache`](_autosummary/mlx.core.metal.clear_cache.html#mlx.core.metal.clear_cache "mlx.core.metal.clear_cache")
() | Clear the memory cache. |
| [`start_capture`](_autosummary/mlx.core.metal.start_capture.html#mlx.core.metal.start_capture "mlx.core.metal.start_capture")
(path) | Start a Metal capture. |
| [`stop_capture`](_autosummary/mlx.core.metal.stop_capture.html#mlx.core.metal.stop_capture "mlx.core.metal.stop_capture")
() | Stop a Metal capture. |
---
# Distributed Communication — MLX 0.24.0 documentation
[Skip to main content](#main-content)
Back to top
Ctrl+K
[](../index.html)
[](https://github.com/ml-explore/mlx "Source repository")
* [.rst](../_sources/python/distributed.rst "Download source file")
* .pdf
Distributed Communication
=========================
Distributed Communication[#](#distributed-communication "Link to this heading")
================================================================================
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.
| | |
| --- | --- |
| [`Group`](_autosummary/mlx.core.distributed.Group.html#mlx.core.distributed.Group "mlx.core.distributed.Group") | An [`mlx.core.distributed.Group`](_autosummary/mlx.core.distributed.Group.html#mlx.core.distributed.Group "mlx.core.distributed.Group")
represents a group of independent mlx processes that can communicate. |
| [`is_available`](_autosummary/mlx.core.distributed.is_available.html#mlx.core.distributed.is_available "mlx.core.distributed.is_available")
() | Check if a communication backend is available. |
| [`init`](_autosummary/mlx.core.distributed.init.html#mlx.core.distributed.init "mlx.core.distributed.init")
(\[strict, backend\]) | Initialize the communication backend and create the global communication group. |
| [`all_sum`](_autosummary/mlx.core.distributed.all_sum.html#mlx.core.distributed.all_sum "mlx.core.distributed.all_sum")
(x, \*\[, group, stream\]) | All reduce sum. |
| [`all_gather`](_autosummary/mlx.core.distributed.all_gather.html#mlx.core.distributed.all_gather "mlx.core.distributed.all_gather")
(x, \*\[, group, stream\]) | Gather arrays from all processes. |
| [`send`](_autosummary/mlx.core.distributed.send.html#mlx.core.distributed.send "mlx.core.distributed.send")
(x, dst, \*\[, group, stream\]) | Send an array from the current process to the process that has rank `dst` in the group. |
| [`recv`](_autosummary/mlx.core.distributed.recv.html#mlx.core.distributed.recv "mlx.core.distributed.recv")
(shape, dtype, src, \*\[, group, stream\]) | Recv an array with shape `shape` and dtype `dtype` from process with rank `src`. |
| [`recv_like`](_autosummary/mlx.core.distributed.recv_like.html#mlx.core.distributed.recv_like "mlx.core.distributed.recv_like")
(x, src, \*\[, group, stream\]) | Recv an array with shape and type like `x` from process with rank `src`. |
---
# Neural Networks — MLX 0.24.0 documentation
[Skip to main content](#main-content)
Back to top
Ctrl+K
[](../index.html)
[](https://github.com/ml-explore/mlx "Source repository")
* [.rst](../_sources/python/nn.rst "Download source file")
* .pdf
Neural Networks
===============
Contents
--------
Neural Networks[#](#neural-networks "Link to this heading")
============================================================
Writing arbitrarily complex neural networks in MLX can be done using only [`mlx.core.array`](_autosummary/mlx.core.array.html#mlx.core.array "mlx.core.array")
and [`mlx.core.value_and_grad()`](_autosummary/mlx.core.value_and_grad.html#mlx.core.value_and_grad "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 `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[#](#quick-start-with-neural-networks "Link to this heading")
----------------------------------------------------------------------------------------------
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)
The Module Class[#](#the-module-class "Link to this heading")
--------------------------------------------------------------
The workhorse of any neural network library is the [`Module`](nn/module.html#mlx.nn.Module "mlx.nn.Module")
class. In MLX the [`Module`](nn/module.html#mlx.nn.Module "mlx.nn.Module")
class is a container of [`mlx.core.array`](_autosummary/mlx.core.array.html#mlx.core.array "mlx.core.array")
or [`Module`](nn/module.html#mlx.nn.Module "mlx.nn.Module")
instances. Its main function is to provide a way to recursively **access** and **update** its parameters and those of its submodules.
### Parameters[#](#parameters "Link to this heading")
A parameter of a module is any public member of type [`mlx.core.array`](_autosummary/mlx.core.array.html#mlx.core.array "mlx.core.array")
(its name should not start with `_`). It can be arbitrarily nested in other [`Module`](nn/module.html#mlx.nn.Module "mlx.nn.Module")
instances or lists and dictionaries.
[`Module.parameters()`](nn/_autosummary/mlx.nn.Module.parameters.html#mlx.nn.Module.parameters "mlx.nn.Module.parameters")
can be used to extract a nested dictionary with all the parameters of a module and its submodules.
A [`Module`](nn/module.html#mlx.nn.Module "mlx.nn.Module")
can also keep track of “frozen” parameters. See the [`Module.freeze()`](nn/_autosummary/mlx.nn.Module.freeze.html#mlx.nn.Module.freeze "mlx.nn.Module.freeze")
method for more details. [`mlx.nn.value_and_grad()`](_autosummary/mlx.nn.value_and_grad.html#mlx.nn.value_and_grad "mlx.nn.value_and_grad")
the gradients returned will be with respect to these trainable parameters.
### Updating the Parameters[#](#updating-the-parameters "Link to this heading")
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 [`Module.update()`](nn/_autosummary/mlx.nn.Module.update.html#mlx.nn.Module.update "mlx.nn.Module.update")
.
### Inspecting Modules[#](#inspecting-modules "Link to this heading")
The simplest way to see the model architecture is to print it. Following along with the above example, you can print the `MLP` with:
print(mlp)
This will display:
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 [`Module`](nn/module.html#mlx.nn.Module "mlx.nn.Module")
you can use [`mlx.utils.tree_map()`](_autosummary/mlx.utils.tree_map.html#mlx.utils.tree_map "mlx.utils.tree_map")
on the parameters. For example, to see the shapes of all the parameters in a [`Module`](nn/module.html#mlx.nn.Module "mlx.nn.Module")
do:
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 [`Module`](nn/module.html#mlx.nn.Module "mlx.nn.Module")
with:
from mlx.utils import tree\_flatten
num\_params \= sum(v.size for \_, v in tree\_flatten(mlp.parameters()))
Value and Grad[#](#value-and-grad "Link to this heading")
----------------------------------------------------------
Using a [`Module`](nn/module.html#mlx.nn.Module "mlx.nn.Module")
does not preclude using MLX’s high order function transformations ([`mlx.core.value_and_grad()`](_autosummary/mlx.core.value_and_grad.html#mlx.core.value_and_grad "mlx.core.value_and_grad")
, [`mlx.core.grad()`](_autosummary/mlx.core.grad.html#mlx.core.grad "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
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, [`mlx.nn.value_and_grad()`](_autosummary/mlx.nn.value_and_grad.html#mlx.nn.value_and_grad "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 [`Module.update()`](nn/_autosummary/mlx.nn.Module.update.html#mlx.nn.Module.update "mlx.nn.Module.update")
to make sure the model is using the provided parameters.
* it calls [`mlx.core.value_and_grad()`](_autosummary/mlx.core.value_and_grad.html#mlx.core.value_and_grad "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 [`mlx.core.value_and_grad()`](_autosummary/mlx.core.value_and_grad.html#mlx.core.value_and_grad "mlx.core.value_and_grad")
| | |
| --- | --- |
| [`value_and_grad`](_autosummary/mlx.nn.value_and_grad.html#mlx.nn.value_and_grad "mlx.nn.value_and_grad")
(model, fn) | Transform the passed function `fn` to a function that computes the gradients of `fn` wrt the model's trainable parameters and also its value. |
| [`quantize`](_autosummary/mlx.nn.quantize.html#mlx.nn.quantize "mlx.nn.quantize")
(model\[, group\_size, bits, ...\]) | Quantize the sub-modules of a module according to a predicate. |
| [`average_gradients`](_autosummary/mlx.nn.average_gradients.html#mlx.nn.average_gradients "mlx.nn.average_gradients")
(gradients\[, group, ...\]) | Average the gradients across the distributed processes in the passed group. |
* [Module](nn/module.html)
* [`Module`](nn/module.html#mlx.nn.Module)
* [mlx.nn.Module.training](nn/_autosummary/mlx.nn.Module.training.html)
* [`Module.training`](nn/_autosummary/mlx.nn.Module.training.html#mlx.nn.Module.training)
* [mlx.nn.Module.state](nn/_autosummary/mlx.nn.Module.state.html)
* [`Module.state`](nn/_autosummary/mlx.nn.Module.state.html#mlx.nn.Module.state)
* [mlx.nn.Module.apply](nn/_autosummary/mlx.nn.Module.apply.html)
* [`Module.apply()`](nn/_autosummary/mlx.nn.Module.apply.html#mlx.nn.Module.apply)
* [mlx.nn.Module.apply\_to\_modules](nn/_autosummary/mlx.nn.Module.apply_to_modules.html)
* [`Module.apply_to_modules()`](nn/_autosummary/mlx.nn.Module.apply_to_modules.html#mlx.nn.Module.apply_to_modules)
* [mlx.nn.Module.children](nn/_autosummary/mlx.nn.Module.children.html)
* [`Module.children()`](nn/_autosummary/mlx.nn.Module.children.html#mlx.nn.Module.children)
* [mlx.nn.Module.eval](nn/_autosummary/mlx.nn.Module.eval.html)
* [`Module.eval()`](nn/_autosummary/mlx.nn.Module.eval.html#mlx.nn.Module.eval)
* [mlx.nn.Module.filter\_and\_map](nn/_autosummary/mlx.nn.Module.filter_and_map.html)
* [`Module.filter_and_map()`](nn/_autosummary/mlx.nn.Module.filter_and_map.html#mlx.nn.Module.filter_and_map)
* [mlx.nn.Module.freeze](nn/_autosummary/mlx.nn.Module.freeze.html)
* [`Module.freeze()`](nn/_autosummary/mlx.nn.Module.freeze.html#mlx.nn.Module.freeze)
* [mlx.nn.Module.leaf\_modules](nn/_autosummary/mlx.nn.Module.leaf_modules.html)
* [`Module.leaf_modules()`](nn/_autosummary/mlx.nn.Module.leaf_modules.html#mlx.nn.Module.leaf_modules)
* [mlx.nn.Module.load\_weights](nn/_autosummary/mlx.nn.Module.load_weights.html)
* [`Module.load_weights()`](nn/_autosummary/mlx.nn.Module.load_weights.html#mlx.nn.Module.load_weights)
* [mlx.nn.Module.modules](nn/_autosummary/mlx.nn.Module.modules.html)
* [`Module.modules()`](nn/_autosummary/mlx.nn.Module.modules.html#mlx.nn.Module.modules)
* [mlx.nn.Module.named\_modules](nn/_autosummary/mlx.nn.Module.named_modules.html)
* [`Module.named_modules()`](nn/_autosummary/mlx.nn.Module.named_modules.html#mlx.nn.Module.named_modules)
* [mlx.nn.Module.parameters](nn/_autosummary/mlx.nn.Module.parameters.html)
* [`Module.parameters()`](nn/_autosummary/mlx.nn.Module.parameters.html#mlx.nn.Module.parameters)
* [mlx.nn.Module.save\_weights](nn/_autosummary/mlx.nn.Module.save_weights.html)
* [`Module.save_weights()`](nn/_autosummary/mlx.nn.Module.save_weights.html#mlx.nn.Module.save_weights)
* [mlx.nn.Module.set\_dtype](nn/_autosummary/mlx.nn.Module.set_dtype.html)
* [`Module.set_dtype()`](nn/_autosummary/mlx.nn.Module.set_dtype.html#mlx.nn.Module.set_dtype)
* [mlx.nn.Module.train](nn/_autosummary/mlx.nn.Module.train.html)
* [`Module.train()`](nn/_autosummary/mlx.nn.Module.train.html#mlx.nn.Module.train)
* [mlx.nn.Module.trainable\_parameters](nn/_autosummary/mlx.nn.Module.trainable_parameters.html)
* [`Module.trainable_parameters()`](nn/_autosummary/mlx.nn.Module.trainable_parameters.html#mlx.nn.Module.trainable_parameters)
* [mlx.nn.Module.unfreeze](nn/_autosummary/mlx.nn.Module.unfreeze.html)
* [`Module.unfreeze()`](nn/_autosummary/mlx.nn.Module.unfreeze.html#mlx.nn.Module.unfreeze)
* [mlx.nn.Module.update](nn/_autosummary/mlx.nn.Module.update.html)
* [`Module.update()`](nn/_autosummary/mlx.nn.Module.update.html#mlx.nn.Module.update)
* [mlx.nn.Module.update\_modules](nn/_autosummary/mlx.nn.Module.update_modules.html)
* [`Module.update_modules()`](nn/_autosummary/mlx.nn.Module.update_modules.html#mlx.nn.Module.update_modules)
* [Layers](nn/layers.html)
* [mlx.nn.ALiBi](nn/_autosummary/mlx.nn.ALiBi.html)
* [`ALiBi`](nn/_autosummary/mlx.nn.ALiBi.html#mlx.nn.ALiBi)
* [mlx.nn.AvgPool1d](nn/_autosummary/mlx.nn.AvgPool1d.html)
* [`AvgPool1d`](nn/_autosummary/mlx.nn.AvgPool1d.html#mlx.nn.AvgPool1d)
* [mlx.nn.AvgPool2d](nn/_autosummary/mlx.nn.AvgPool2d.html)
* [`AvgPool2d`](nn/_autosummary/mlx.nn.AvgPool2d.html#mlx.nn.AvgPool2d)
* [mlx.nn.AvgPool3d](nn/_autosummary/mlx.nn.AvgPool3d.html)
* [`AvgPool3d`](nn/_autosummary/mlx.nn.AvgPool3d.html#mlx.nn.AvgPool3d)
* [mlx.nn.BatchNorm](nn/_autosummary/mlx.nn.BatchNorm.html)
* [`BatchNorm`](nn/_autosummary/mlx.nn.BatchNorm.html#mlx.nn.BatchNorm)
* [mlx.nn.CELU](nn/_autosummary/mlx.nn.CELU.html)
* [`CELU`](nn/_autosummary/mlx.nn.CELU.html#mlx.nn.CELU)
* [mlx.nn.Conv1d](nn/_autosummary/mlx.nn.Conv1d.html)
* [`Conv1d`](nn/_autosummary/mlx.nn.Conv1d.html#mlx.nn.Conv1d)
* [mlx.nn.Conv2d](nn/_autosummary/mlx.nn.Conv2d.html)
* [`Conv2d`](nn/_autosummary/mlx.nn.Conv2d.html#mlx.nn.Conv2d)
* [mlx.nn.Conv3d](nn/_autosummary/mlx.nn.Conv3d.html)
* [`Conv3d`](nn/_autosummary/mlx.nn.Conv3d.html#mlx.nn.Conv3d)
* [mlx.nn.ConvTranspose1d](nn/_autosummary/mlx.nn.ConvTranspose1d.html)
* [`ConvTranspose1d`](nn/_autosummary/mlx.nn.ConvTranspose1d.html#mlx.nn.ConvTranspose1d)
* [mlx.nn.ConvTranspose2d](nn/_autosummary/mlx.nn.ConvTranspose2d.html)
* [`ConvTranspose2d`](nn/_autosummary/mlx.nn.ConvTranspose2d.html#mlx.nn.ConvTranspose2d)
* [mlx.nn.ConvTranspose3d](nn/_autosummary/mlx.nn.ConvTranspose3d.html)
* [`ConvTranspose3d`](nn/_autosummary/mlx.nn.ConvTranspose3d.html#mlx.nn.ConvTranspose3d)
* [mlx.nn.Dropout](nn/_autosummary/mlx.nn.Dropout.html)
* [`Dropout`](nn/_autosummary/mlx.nn.Dropout.html#mlx.nn.Dropout)
* [mlx.nn.Dropout2d](nn/_autosummary/mlx.nn.Dropout2d.html)
* [`Dropout2d`](nn/_autosummary/mlx.nn.Dropout2d.html#mlx.nn.Dropout2d)
* [mlx.nn.Dropout3d](nn/_autosummary/mlx.nn.Dropout3d.html)
* [`Dropout3d`](nn/_autosummary/mlx.nn.Dropout3d.html#mlx.nn.Dropout3d)
* [mlx.nn.Embedding](nn/_autosummary/mlx.nn.Embedding.html)
* [`Embedding`](nn/_autosummary/mlx.nn.Embedding.html#mlx.nn.Embedding)
* [mlx.nn.ELU](nn/_autosummary/mlx.nn.ELU.html)
* [`ELU`](nn/_autosummary/mlx.nn.ELU.html#mlx.nn.ELU)
* [mlx.nn.GELU](nn/_autosummary/mlx.nn.GELU.html)
* [`GELU`](nn/_autosummary/mlx.nn.GELU.html#mlx.nn.GELU)
* [mlx.nn.GLU](nn/_autosummary/mlx.nn.GLU.html)
* [`GLU`](nn/_autosummary/mlx.nn.GLU.html#mlx.nn.GLU)
* [mlx.nn.GroupNorm](nn/_autosummary/mlx.nn.GroupNorm.html)
* [`GroupNorm`](nn/_autosummary/mlx.nn.GroupNorm.html#mlx.nn.GroupNorm)
* [mlx.nn.GRU](nn/_autosummary/mlx.nn.GRU.html)
* [`GRU`](nn/_autosummary/mlx.nn.GRU.html#mlx.nn.GRU)
* [mlx.nn.HardShrink](nn/_autosummary/mlx.nn.HardShrink.html)
* [`HardShrink`](nn/_autosummary/mlx.nn.HardShrink.html#mlx.nn.HardShrink)
* [mlx.nn.HardTanh](nn/_autosummary/mlx.nn.HardTanh.html)
* [`HardTanh`](nn/_autosummary/mlx.nn.HardTanh.html#mlx.nn.HardTanh)
* [mlx.nn.Hardswish](nn/_autosummary/mlx.nn.Hardswish.html)
* [`Hardswish`](nn/_autosummary/mlx.nn.Hardswish.html#mlx.nn.Hardswish)
* [mlx.nn.InstanceNorm](nn/_autosummary/mlx.nn.InstanceNorm.html)
* [`InstanceNorm`](nn/_autosummary/mlx.nn.InstanceNorm.html#mlx.nn.InstanceNorm)
* [mlx.nn.LayerNorm](nn/_autosummary/mlx.nn.LayerNorm.html)
* [`LayerNorm`](nn/_autosummary/mlx.nn.LayerNorm.html#mlx.nn.LayerNorm)
* [mlx.nn.LeakyReLU](nn/_autosummary/mlx.nn.LeakyReLU.html)
* [`LeakyReLU`](nn/_autosummary/mlx.nn.LeakyReLU.html#mlx.nn.LeakyReLU)
* [mlx.nn.Linear](nn/_autosummary/mlx.nn.Linear.html)
* [`Linear`](nn/_autosummary/mlx.nn.Linear.html#mlx.nn.Linear)
* [mlx.nn.LogSigmoid](nn/_autosummary/mlx.nn.LogSigmoid.html)
* [`LogSigmoid`](nn/_autosummary/mlx.nn.LogSigmoid.html#mlx.nn.LogSigmoid)
* [mlx.nn.LogSoftmax](nn/_autosummary/mlx.nn.LogSoftmax.html)
* [`LogSoftmax`](nn/_autosummary/mlx.nn.LogSoftmax.html#mlx.nn.LogSoftmax)
* [mlx.nn.LSTM](nn/_autosummary/mlx.nn.LSTM.html)
* [`LSTM`](nn/_autosummary/mlx.nn.LSTM.html#mlx.nn.LSTM)
* [mlx.nn.MaxPool1d](nn/_autosummary/mlx.nn.MaxPool1d.html)
* [`MaxPool1d`](nn/_autosummary/mlx.nn.MaxPool1d.html#mlx.nn.MaxPool1d)
* [mlx.nn.MaxPool2d](nn/_autosummary/mlx.nn.MaxPool2d.html)
* [`MaxPool2d`](nn/_autosummary/mlx.nn.MaxPool2d.html#mlx.nn.MaxPool2d)
* [mlx.nn.MaxPool3d](nn/_autosummary/mlx.nn.MaxPool3d.html)
* [`MaxPool3d`](nn/_autosummary/mlx.nn.MaxPool3d.html#mlx.nn.MaxPool3d)
* [mlx.nn.Mish](nn/_autosummary/mlx.nn.Mish.html)
* [`Mish`](nn/_autosummary/mlx.nn.Mish.html#mlx.nn.Mish)
* [mlx.nn.MultiHeadAttention](nn/_autosummary/mlx.nn.MultiHeadAttention.html)
* [`MultiHeadAttention`](nn/_autosummary/mlx.nn.MultiHeadAttention.html#mlx.nn.MultiHeadAttention)
* [mlx.nn.PReLU](nn/_autosummary/mlx.nn.PReLU.html)
* [`PReLU`](nn/_autosummary/mlx.nn.PReLU.html#mlx.nn.PReLU)
* [mlx.nn.QuantizedEmbedding](nn/_autosummary/mlx.nn.QuantizedEmbedding.html)
* [`QuantizedEmbedding`](nn/_autosummary/mlx.nn.QuantizedEmbedding.html#mlx.nn.QuantizedEmbedding)
* [mlx.nn.QuantizedLinear](nn/_autosummary/mlx.nn.QuantizedLinear.html)
* [`QuantizedLinear`](nn/_autosummary/mlx.nn.QuantizedLinear.html#mlx.nn.QuantizedLinear)
* [mlx.nn.RMSNorm](nn/_autosummary/mlx.nn.RMSNorm.html)
* [`RMSNorm`](nn/_autosummary/mlx.nn.RMSNorm.html#mlx.nn.RMSNorm)
* [mlx.nn.ReLU](nn/_autosummary/mlx.nn.ReLU.html)
* [`ReLU`](nn/_autosummary/mlx.nn.ReLU.html#mlx.nn.ReLU)
* [mlx.nn.ReLU6](nn/_autosummary/mlx.nn.ReLU6.html)
* [`ReLU6`](nn/_autosummary/mlx.nn.ReLU6.html#mlx.nn.ReLU6)
* [mlx.nn.RNN](nn/_autosummary/mlx.nn.RNN.html)
* [`RNN`](nn/_autosummary/mlx.nn.RNN.html#mlx.nn.RNN)
* [mlx.nn.RoPE](nn/_autosummary/mlx.nn.RoPE.html)
* [`RoPE`](nn/_autosummary/mlx.nn.RoPE.html#mlx.nn.RoPE)
* [mlx.nn.SELU](nn/_autosummary/mlx.nn.SELU.html)
* [`SELU`](nn/_autosummary/mlx.nn.SELU.html#mlx.nn.SELU)
* [mlx.nn.Sequential](nn/_autosummary/mlx.nn.Sequential.html)
* [`Sequential`](nn/_autosummary/mlx.nn.Sequential.html#mlx.nn.Sequential)
* [mlx.nn.Sigmoid](nn/_autosummary/mlx.nn.Sigmoid.html)
* [`Sigmoid`](nn/_autosummary/mlx.nn.Sigmoid.html#mlx.nn.Sigmoid)
* [mlx.nn.SiLU](nn/_autosummary/mlx.nn.SiLU.html)
* [`SiLU`](nn/_autosummary/mlx.nn.SiLU.html#mlx.nn.SiLU)
* [mlx.nn.SinusoidalPositionalEncoding](nn/_autosummary/mlx.nn.SinusoidalPositionalEncoding.html)
* [`SinusoidalPositionalEncoding`](nn/_autosummary/mlx.nn.SinusoidalPositionalEncoding.html#mlx.nn.SinusoidalPositionalEncoding)
* [mlx.nn.Softmin](nn/_autosummary/mlx.nn.Softmin.html)
* [`Softmin`](nn/_autosummary/mlx.nn.Softmin.html#mlx.nn.Softmin)
* [mlx.nn.Softshrink](nn/_autosummary/mlx.nn.Softshrink.html)
* [`Softshrink`](nn/_autosummary/mlx.nn.Softshrink.html#mlx.nn.Softshrink)
* [mlx.nn.Softsign](nn/_autosummary/mlx.nn.Softsign.html)
* [`Softsign`](nn/_autosummary/mlx.nn.Softsign.html#mlx.nn.Softsign)
* [mlx.nn.Softmax](nn/_autosummary/mlx.nn.Softmax.html)
* [`Softmax`](nn/_autosummary/mlx.nn.Softmax.html#mlx.nn.Softmax)
* [mlx.nn.Softplus](nn/_autosummary/mlx.nn.Softplus.html)
* [`Softplus`](nn/_autosummary/mlx.nn.Softplus.html#mlx.nn.Softplus)
* [mlx.nn.Step](nn/_autosummary/mlx.nn.Step.html)
* [`Step`](nn/_autosummary/mlx.nn.Step.html#mlx.nn.Step)
* [mlx.nn.Tanh](nn/_autosummary/mlx.nn.Tanh.html)
* [`Tanh`](nn/_autosummary/mlx.nn.Tanh.html#mlx.nn.Tanh)
* [mlx.nn.Transformer](nn/_autosummary/mlx.nn.Transformer.html)
* [`Transformer`](nn/_autosummary/mlx.nn.Transformer.html#mlx.nn.Transformer)
* [mlx.nn.Upsample](nn/_autosummary/mlx.nn.Upsample.html)
* [`Upsample`](nn/_autosummary/mlx.nn.Upsample.html#mlx.nn.Upsample)
* [Functions](nn/functions.html)
* [mlx.nn.elu](nn/_autosummary_functions/mlx.nn.elu.html)
* [`elu`](nn/_autosummary_functions/mlx.nn.elu.html#mlx.nn.elu)
* [mlx.nn.celu](nn/_autosummary_functions/mlx.nn.celu.html)
* [`celu`](nn/_autosummary_functions/mlx.nn.celu.html#mlx.nn.celu)
* [mlx.nn.gelu](nn/_autosummary_functions/mlx.nn.gelu.html)
* [`gelu`](nn/_autosummary_functions/mlx.nn.gelu.html#mlx.nn.gelu)
* [mlx.nn.gelu\_approx](nn/_autosummary_functions/mlx.nn.gelu_approx.html)
* [`gelu_approx`](nn/_autosummary_functions/mlx.nn.gelu_approx.html#mlx.nn.gelu_approx)
* [mlx.nn.gelu\_fast\_approx](nn/_autosummary_functions/mlx.nn.gelu_fast_approx.html)
* [`gelu_fast_approx`](nn/_autosummary_functions/mlx.nn.gelu_fast_approx.html#mlx.nn.gelu_fast_approx)
* [mlx.nn.glu](nn/_autosummary_functions/mlx.nn.glu.html)
* [`glu`](nn/_autosummary_functions/mlx.nn.glu.html#mlx.nn.glu)
* [mlx.nn.hard\_shrink](nn/_autosummary_functions/mlx.nn.hard_shrink.html)
* [`hard_shrink`](nn/_autosummary_functions/mlx.nn.hard_shrink.html#mlx.nn.hard_shrink)
* [mlx.nn.hard\_tanh](nn/_autosummary_functions/mlx.nn.hard_tanh.html)
* [`hard_tanh`](nn/_autosummary_functions/mlx.nn.hard_tanh.html#mlx.nn.hard_tanh)
* [mlx.nn.hardswish](nn/_autosummary_functions/mlx.nn.hardswish.html)
* [`hardswish`](nn/_autosummary_functions/mlx.nn.hardswish.html#mlx.nn.hardswish)
* [mlx.nn.leaky\_relu](nn/_autosummary_functions/mlx.nn.leaky_relu.html)
* [`leaky_relu`](nn/_autosummary_functions/mlx.nn.leaky_relu.html#mlx.nn.leaky_relu)
* [mlx.nn.log\_sigmoid](nn/_autosummary_functions/mlx.nn.log_sigmoid.html)
* [`log_sigmoid`](nn/_autosummary_functions/mlx.nn.log_sigmoid.html#mlx.nn.log_sigmoid)
* [mlx.nn.log\_softmax](nn/_autosummary_functions/mlx.nn.log_softmax.html)
* [`log_softmax`](nn/_autosummary_functions/mlx.nn.log_softmax.html#mlx.nn.log_softmax)
* [mlx.nn.mish](nn/_autosummary_functions/mlx.nn.mish.html)
* [`mish`](nn/_autosummary_functions/mlx.nn.mish.html#mlx.nn.mish)
* [mlx.nn.prelu](nn/_autosummary_functions/mlx.nn.prelu.html)
* [`prelu`](nn/_autosummary_functions/mlx.nn.prelu.html#mlx.nn.prelu)
* [mlx.nn.relu](nn/_autosummary_functions/mlx.nn.relu.html)
* [`relu`](nn/_autosummary_functions/mlx.nn.relu.html#mlx.nn.relu)
* [mlx.nn.relu6](nn/_autosummary_functions/mlx.nn.relu6.html)
* [`relu6`](nn/_autosummary_functions/mlx.nn.relu6.html#mlx.nn.relu6)
* [mlx.nn.selu](nn/_autosummary_functions/mlx.nn.selu.html)
* [`selu`](nn/_autosummary_functions/mlx.nn.selu.html#mlx.nn.selu)
* [mlx.nn.sigmoid](nn/_autosummary_functions/mlx.nn.sigmoid.html)
* [`sigmoid`](nn/_autosummary_functions/mlx.nn.sigmoid.html#mlx.nn.sigmoid)
* [mlx.nn.silu](nn/_autosummary_functions/mlx.nn.silu.html)
* [`silu`](nn/_autosummary_functions/mlx.nn.silu.html#mlx.nn.silu)
* [mlx.nn.softmax](nn/_autosummary_functions/mlx.nn.softmax.html)
* [`softmax`](nn/_autosummary_functions/mlx.nn.softmax.html#mlx.nn.softmax)
* [mlx.nn.softmin](nn/_autosummary_functions/mlx.nn.softmin.html)
* [`softmin`](nn/_autosummary_functions/mlx.nn.softmin.html#mlx.nn.softmin)
* [mlx.nn.softplus](nn/_autosummary_functions/mlx.nn.softplus.html)
* [`softplus`](nn/_autosummary_functions/mlx.nn.softplus.html#mlx.nn.softplus)
* [mlx.nn.softshrink](nn/_autosummary_functions/mlx.nn.softshrink.html)
* [`softshrink`](nn/_autosummary_functions/mlx.nn.softshrink.html#mlx.nn.softshrink)
* [mlx.nn.step](nn/_autosummary_functions/mlx.nn.step.html)
* [`step`](nn/_autosummary_functions/mlx.nn.step.html#mlx.nn.step)
* [mlx.nn.tanh](nn/_autosummary_functions/mlx.nn.tanh.html)
* [`tanh`](nn/_autosummary_functions/mlx.nn.tanh.html#mlx.nn.tanh)
* [Loss Functions](nn/losses.html)
* [mlx.nn.losses.binary\_cross\_entropy](nn/_autosummary_functions/mlx.nn.losses.binary_cross_entropy.html)
* [`binary_cross_entropy`](nn/_autosummary_functions/mlx.nn.losses.binary_cross_entropy.html#mlx.nn.losses.binary_cross_entropy)
* [mlx.nn.losses.cosine\_similarity\_loss](nn/_autosummary_functions/mlx.nn.losses.cosine_similarity_loss.html)
* [`cosine_similarity_loss`](nn/_autosummary_functions/mlx.nn.losses.cosine_similarity_loss.html#mlx.nn.losses.cosine_similarity_loss)
* [mlx.nn.losses.cross\_entropy](nn/_autosummary_functions/mlx.nn.losses.cross_entropy.html)
* [`cross_entropy`](nn/_autosummary_functions/mlx.nn.losses.cross_entropy.html#mlx.nn.losses.cross_entropy)
* [mlx.nn.losses.gaussian\_nll\_loss](nn/_autosummary_functions/mlx.nn.losses.gaussian_nll_loss.html)
* [`gaussian_nll_loss`](nn/_autosummary_functions/mlx.nn.losses.gaussian_nll_loss.html#mlx.nn.losses.gaussian_nll_loss)
* [mlx.nn.losses.hinge\_loss](nn/_autosummary_functions/mlx.nn.losses.hinge_loss.html)
* [`hinge_loss`](nn/_autosummary_functions/mlx.nn.losses.hinge_loss.html#mlx.nn.losses.hinge_loss)
* [mlx.nn.losses.huber\_loss](nn/_autosummary_functions/mlx.nn.losses.huber_loss.html)
* [`huber_loss`](nn/_autosummary_functions/mlx.nn.losses.huber_loss.html#mlx.nn.losses.huber_loss)
* [mlx.nn.losses.kl\_div\_loss](nn/_autosummary_functions/mlx.nn.losses.kl_div_loss.html)
* [`kl_div_loss`](nn/_autosummary_functions/mlx.nn.losses.kl_div_loss.html#mlx.nn.losses.kl_div_loss)
* [mlx.nn.losses.l1\_loss](nn/_autosummary_functions/mlx.nn.losses.l1_loss.html)
* [`l1_loss`](nn/_autosummary_functions/mlx.nn.losses.l1_loss.html#mlx.nn.losses.l1_loss)
* [mlx.nn.losses.log\_cosh\_loss](nn/_autosummary_functions/mlx.nn.losses.log_cosh_loss.html)
* [`log_cosh_loss`](nn/_autosummary_functions/mlx.nn.losses.log_cosh_loss.html#mlx.nn.losses.log_cosh_loss)
* [mlx.nn.losses.margin\_ranking\_loss](nn/_autosummary_functions/mlx.nn.losses.margin_ranking_loss.html)
* [`margin_ranking_loss`](nn/_autosummary_functions/mlx.nn.losses.margin_ranking_loss.html#mlx.nn.losses.margin_ranking_loss)
* [mlx.nn.losses.mse\_loss](nn/_autosummary_functions/mlx.nn.losses.mse_loss.html)
* [`mse_loss`](nn/_autosummary_functions/mlx.nn.losses.mse_loss.html#mlx.nn.losses.mse_loss)
* [mlx.nn.losses.nll\_loss](nn/_autosummary_functions/mlx.nn.losses.nll_loss.html)
* [`nll_loss`](nn/_autosummary_functions/mlx.nn.losses.nll_loss.html#mlx.nn.losses.nll_loss)
* [mlx.nn.losses.smooth\_l1\_loss](nn/_autosummary_functions/mlx.nn.losses.smooth_l1_loss.html)
* [`smooth_l1_loss`](nn/_autosummary_functions/mlx.nn.losses.smooth_l1_loss.html#mlx.nn.losses.smooth_l1_loss)
* [mlx.nn.losses.triplet\_loss](nn/_autosummary_functions/mlx.nn.losses.triplet_loss.html)
* [`triplet_loss`](nn/_autosummary_functions/mlx.nn.losses.triplet_loss.html#mlx.nn.losses.triplet_loss)
* [Initializers](nn/init.html)
* [mlx.nn.init.constant](nn/_autosummary/mlx.nn.init.constant.html)
* [`constant()`](nn/_autosummary/mlx.nn.init.constant.html#mlx.nn.init.constant)
* [mlx.nn.init.normal](nn/_autosummary/mlx.nn.init.normal.html)
* [`normal()`](nn/_autosummary/mlx.nn.init.normal.html#mlx.nn.init.normal)
* [mlx.nn.init.uniform](nn/_autosummary/mlx.nn.init.uniform.html)
* [`uniform()`](nn/_autosummary/mlx.nn.init.uniform.html#mlx.nn.init.uniform)
* [mlx.nn.init.identity](nn/_autosummary/mlx.nn.init.identity.html)
* [`identity()`](nn/_autosummary/mlx.nn.init.identity.html#mlx.nn.init.identity)
* [mlx.nn.init.glorot\_normal](nn/_autosummary/mlx.nn.init.glorot_normal.html)
* [`glorot_normal()`](nn/_autosummary/mlx.nn.init.glorot_normal.html#mlx.nn.init.glorot_normal)
* [mlx.nn.init.glorot\_uniform](nn/_autosummary/mlx.nn.init.glorot_uniform.html)
* [`glorot_uniform()`](nn/_autosummary/mlx.nn.init.glorot_uniform.html#mlx.nn.init.glorot_uniform)
* [mlx.nn.init.he\_normal](nn/_autosummary/mlx.nn.init.he_normal.html)
* [`he_normal()`](nn/_autosummary/mlx.nn.init.he_normal.html#mlx.nn.init.he_normal)
* [mlx.nn.init.he\_uniform](nn/_autosummary/mlx.nn.init.he_uniform.html)
* [`he_uniform()`](nn/_autosummary/mlx.nn.init.he_uniform.html#mlx.nn.init.he_uniform)
Contents
---
# Optimizers — MLX 0.24.0 documentation
[Skip to main content](#main-content)
Back to top
Ctrl+K
[](../index.html)
[](https://github.com/ml-explore/mlx "Source repository")
* [.rst](../_sources/python/optimizers.rst "Download source file")
* .pdf
Optimizers
==========
Contents
--------
Optimizers[#](#optimizers "Link to this heading")
==================================================
The optimizers in MLX can be used both with `mlx.nn` but also with pure `mlx.core` functions. A typical example involves calling [`Optimizer.update()`](optimizers/_autosummary/mlx.optimizers.Optimizer.update.html#mlx.optimizers.Optimizer.update "mlx.optimizers.Optimizer.update")
to update a model’s parameters based on the loss gradients and subsequently calling [`mlx.core.eval()`](_autosummary/mlx.core.eval.html#mlx.core.eval "mlx.core.eval")
to evaluate both the model’s parameters and the **optimizer state**.
\# 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[#](#saving-and-loading "Link to this heading")
------------------------------------------------------------------
To serialize an optimizer, save its state. To load an optimizer, load and set the saved state. Here’s a simple example:
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)
mx.save\_safetensors("optimizer.safetensors", dict(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(list(mx.load("optimizer.safetensors").items()))
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.
* [Optimizer](optimizers/optimizer.html)
* [`Optimizer`](optimizers/optimizer.html#mlx.optimizers.Optimizer)
* [mlx.optimizers.Optimizer.state](optimizers/_autosummary/mlx.optimizers.Optimizer.state.html)
* [`Optimizer.state`](optimizers/_autosummary/mlx.optimizers.Optimizer.state.html#mlx.optimizers.Optimizer.state)
* [mlx.optimizers.Optimizer.apply\_gradients](optimizers/_autosummary/mlx.optimizers.Optimizer.apply_gradients.html)
* [`Optimizer.apply_gradients()`](optimizers/_autosummary/mlx.optimizers.Optimizer.apply_gradients.html#mlx.optimizers.Optimizer.apply_gradients)
* [mlx.optimizers.Optimizer.init](optimizers/_autosummary/mlx.optimizers.Optimizer.init.html)
* [`Optimizer.init()`](optimizers/_autosummary/mlx.optimizers.Optimizer.init.html#mlx.optimizers.Optimizer.init)
* [mlx.optimizers.Optimizer.update](optimizers/_autosummary/mlx.optimizers.Optimizer.update.html)
* [`Optimizer.update()`](optimizers/_autosummary/mlx.optimizers.Optimizer.update.html#mlx.optimizers.Optimizer.update)
* [Common Optimizers](optimizers/common_optimizers.html)
* [mlx.optimizers.SGD](optimizers/_autosummary/mlx.optimizers.SGD.html)
* [`SGD`](optimizers/_autosummary/mlx.optimizers.SGD.html#mlx.optimizers.SGD)
* [mlx.optimizers.RMSprop](optimizers/_autosummary/mlx.optimizers.RMSprop.html)
* [`RMSprop`](optimizers/_autosummary/mlx.optimizers.RMSprop.html#mlx.optimizers.RMSprop)
* [mlx.optimizers.Adagrad](optimizers/_autosummary/mlx.optimizers.Adagrad.html)
* [`Adagrad`](optimizers/_autosummary/mlx.optimizers.Adagrad.html#mlx.optimizers.Adagrad)
* [mlx.optimizers.Adafactor](optimizers/_autosummary/mlx.optimizers.Adafactor.html)
* [`Adafactor`](optimizers/_autosummary/mlx.optimizers.Adafactor.html#mlx.optimizers.Adafactor)
* [mlx.optimizers.AdaDelta](optimizers/_autosummary/mlx.optimizers.AdaDelta.html)
* [`AdaDelta`](optimizers/_autosummary/mlx.optimizers.AdaDelta.html#mlx.optimizers.AdaDelta)
* [mlx.optimizers.Adam](optimizers/_autosummary/mlx.optimizers.Adam.html)
* [`Adam`](optimizers/_autosummary/mlx.optimizers.Adam.html#mlx.optimizers.Adam)
* [mlx.optimizers.AdamW](optimizers/_autosummary/mlx.optimizers.AdamW.html)
* [`AdamW`](optimizers/_autosummary/mlx.optimizers.AdamW.html#mlx.optimizers.AdamW)
* [mlx.optimizers.Adamax](optimizers/_autosummary/mlx.optimizers.Adamax.html)
* [`Adamax`](optimizers/_autosummary/mlx.optimizers.Adamax.html#mlx.optimizers.Adamax)
* [mlx.optimizers.Lion](optimizers/_autosummary/mlx.optimizers.Lion.html)
* [`Lion`](optimizers/_autosummary/mlx.optimizers.Lion.html#mlx.optimizers.Lion)
* [Schedulers](optimizers/schedulers.html)
* [mlx.optimizers.cosine\_decay](optimizers/_autosummary/mlx.optimizers.cosine_decay.html)
* [`cosine_decay()`](optimizers/_autosummary/mlx.optimizers.cosine_decay.html#mlx.optimizers.cosine_decay)
* [mlx.optimizers.exponential\_decay](optimizers/_autosummary/mlx.optimizers.exponential_decay.html)
* [`exponential_decay()`](optimizers/_autosummary/mlx.optimizers.exponential_decay.html#mlx.optimizers.exponential_decay)
* [mlx.optimizers.join\_schedules](optimizers/_autosummary/mlx.optimizers.join_schedules.html)
* [`join_schedules()`](optimizers/_autosummary/mlx.optimizers.join_schedules.html#mlx.optimizers.join_schedules)
* [mlx.optimizers.linear\_schedule](optimizers/_autosummary/mlx.optimizers.linear_schedule.html)
* [`linear_schedule()`](optimizers/_autosummary/mlx.optimizers.linear_schedule.html#mlx.optimizers.linear_schedule)
* [mlx.optimizers.step\_decay](optimizers/_autosummary/mlx.optimizers.step_decay.html)
* [`step_decay()`](optimizers/_autosummary/mlx.optimizers.step_decay.html#mlx.optimizers.step_decay)
| | |
| --- | --- |
| [`clip_grad_norm`](_autosummary/mlx.optimizers.clip_grad_norm.html#mlx.optimizers.clip_grad_norm "mlx.optimizers.clip_grad_norm")
(grads, max\_norm) | Clips the global norm of the gradients. |
Contents
---
# Tree Utils — MLX 0.24.0 documentation
[Skip to main content](#main-content)
Back to top
Ctrl+K
[](../index.html)
[](https://github.com/ml-explore/mlx "Source repository")
* [.rst](../_sources/python/tree_utils.rst "Download source file")
* .pdf
Tree Utils
==========
Tree Utils[#](#tree-utils "Link to this heading")
==================================================
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.
| | |
| --- | --- |
| [`tree_flatten`](_autosummary/mlx.utils.tree_flatten.html#mlx.utils.tree_flatten "mlx.utils.tree_flatten")
(tree\[, prefix, is\_leaf\]) | Flattens a Python tree to a list of key, value tuples. |
| [`tree_unflatten`](_autosummary/mlx.utils.tree_unflatten.html#mlx.utils.tree_unflatten "mlx.utils.tree_unflatten")
(tree) | Recreate a Python tree from its flat representation. |
| [`tree_map`](_autosummary/mlx.utils.tree_map.html#mlx.utils.tree_map "mlx.utils.tree_map")
(fn, tree, \*rest\[, is\_leaf\]) | Applies `fn` to the leaves of the Python tree `tree` and returns a new collection with the results. |
| [`tree_map_with_path`](_autosummary/mlx.utils.tree_map_with_path.html#mlx.utils.tree_map_with_path "mlx.utils.tree_map_with_path")
(fn, tree, \*rest\[, ...\]) | Applies `fn` to the path and leaves of the Python tree `tree` and returns a new collection with the results. |
| [`tree_reduce`](_autosummary/mlx.utils.tree_reduce.html#mlx.utils.tree_reduce "mlx.utils.tree_reduce")
(fn, tree\[, initializer, is\_leaf\]) | Applies a reduction to the leaves of a Python tree. |
---
# Metal Debugger — MLX 0.24.0 documentation
[Skip to main content](#main-content)
Back to top
Ctrl+K
[](../index.html)
[](https://github.com/ml-explore/mlx "Source repository")
* [.rst](../_sources/dev/metal_debugger.rst "Download source file")
* .pdf
Metal Debugger
==============
Contents
--------
Metal Debugger[#](#metal-debugger "Link to this heading")
==========================================================
Profiling is a key step for performance optimization. You can build MLX with the `MLX_METAL_DEBUG` option to improve the Metal debugging and optimization workflow. The `MLX_METAL_DEBUG` debug option:
* Records source during Metal compilation, for later inspection while debugging.
* Labels Metal objects such as command queues, improving capture readability.
To build with debugging enabled in Python prepend `CMAKE_ARGS="-DMLX_METAL_DEBUG=ON"` to the build call.
The [`metal.start_capture()`](../python/_autosummary/mlx.core.metal.start_capture.html#mlx.core.metal.start_capture "mlx.core.metal.start_capture")
function initiates a capture of all MLX GPU work.
Note
To capture a GPU trace you must run the application with `MTL_CAPTURE_ENABLED=1`.
import mlx.core as mx
a \= mx.random.uniform(shape\=(512, 512))
b \= mx.random.uniform(shape\=(512, 512))
mx.eval(a, b)
trace\_file \= "mlx\_trace.gputrace"
\# Make sure to run with MTL\_CAPTURE\_ENABLED=1 and
\# that the path trace\_file does not already exist.
mx.metal.start\_capture(trace\_file)
for \_ in range(10):
mx.eval(mx.add(a, b))
mx.metal.stop\_capture()
You can open and replay the GPU trace in Xcode. The `Dependencies` view has a great overview of all operations. Checkout the [Metal debugger documentation](https://developer.apple.com/documentation/xcode/metal-debugger)
for more information.
Xcode Workflow[#](#xcode-workflow "Link to this heading")
----------------------------------------------------------
You can skip saving to a path by running within Xcode. First, generate an Xcode project using CMake.
mkdir build && cd build
cmake .. \-DMLX\_METAL\_DEBUG\=ON \-G Xcode
open mlx.xcodeproj
Select the `metal_capture` example schema and run.
Contents
---
# Custom Metal Kernels — MLX 0.24.0 documentation
[Skip to main content](#main-content)
Back to top
Ctrl+K
[](../index.html)
[](https://github.com/ml-explore/mlx "Source repository")
* [.rst](../_sources/dev/custom_metal_kernels.rst "Download source file")
* .pdf
Custom Metal Kernels
====================
Contents
--------
Custom Metal Kernels[#](#custom-metal-kernels "Link to this heading")
======================================================================
MLX supports writing custom Metal kernels through the Python and C++ APIs.
Simple Example[#](#simple-example "Link to this heading")
----------------------------------------------------------
Let’s write a custom kernel that computes `exp` elementwise:
def exp\_elementwise(a: mx.array):
source \= """
uint elem = thread\_position\_in\_grid.x;
T tmp = inp\[elem\];
out\[elem\] = metal::exp(tmp);
"""
kernel \= mx.fast.metal\_kernel(
name\="myexp",
input\_names\=\["inp"\],
output\_names\=\["out"\],
source\=source,
)
outputs \= kernel(
inputs\=\[a\],
template\=\[("T", mx.float32)\],
grid\=(a.size, 1, 1),
threadgroup\=(256, 1, 1),
output\_shapes\=\[a.shape\],
output\_dtypes\=\[a.dtype\],
)
return outputs\[0\]
a \= mx.random.normal(shape\=(4, 16)).astype(mx.float16)
b \= exp\_elementwise(a)
assert mx.allclose(b, mx.exp(a))
Note
We are only required to pass the body of the Metal kernel in `source`.
The full function signature will be generated using:
* The shapes/dtypes of `inputs`
In the above, `a` is an `mx.array` of type `mx.float16` and we pass it with the key `inp` so we will add `const device float16_t* inp` to the signature. `inp_shape`, `inp_strides` and `inp_ndim` are also added for convenience if they are present in `source`.
* The list of `output_dtypes`
In the above, `out` is an `mx.array` of type `mx.float16` so we add `device float16_t* out`.
* Template parameters passed using `template`
In the above, `template=[("T", mx.float32)]` adds a template of `template ` to the function and instantiates the template with `custom_kernel_myexp_float`. Template parameters can be `mx.core.Dtype`, `int` or `bool`.
* Metal attributes used in `source` such as `[[thread_position_in_grid]]`
These will be added as function arguments. All the attributes defined in Table 5.8 of the [Metal Shading Language Specification](https://developer.apple.com/metal/Metal-Shading-Language-Specification.pdf)
are supported.
Putting this all together, the generated function signature for `myexp` is as follows:
template
\[\[kernel\]\] void custom\_kernel\_myexp\_float(
const device float16\_t\* inp \[\[buffer(0)\]\],
device float16\_t\* out \[\[buffer(1)\]\],
uint3 thread\_position\_in\_grid \[\[thread\_position\_in\_grid\]\]) {
uint elem \= thread\_position\_in\_grid.x;
T tmp \= inp\[elem\];
out\[elem\] \= metal::exp(tmp);
}
template \[\[host\_name("custom\_kernel\_myexp\_float")\]\] \[\[kernel\]\] decltype(custom\_kernel\_myexp\_float) custom\_kernel\_myexp\_float;
Note: `grid` and `threadgroup` are parameters to the Metal [dispatchThreads](https://developer.apple.com/documentation/metal/mtlcomputecommandencoder/2866532-dispatchthreads)
function. This means we will launch `mx.prod(grid)` threads, subdivided into `threadgroup` size threadgroups. For optimal performance, each thread group dimension should be less than or equal to the corresponding grid dimension.
Passing `verbose=True` to `mx.fast.metal_kernel.__call__` will print the generated code for debugging purposes.
Using Shape/Strides[#](#using-shape-strides "Link to this heading")
--------------------------------------------------------------------
`mx.fast.metal_kernel` supports an argument `ensure_row_contiguous` which is `True` by default. This will copy the `mx.array` inputs if needed before the kernel is launched to ensure that the memory layout is row contiguous. Generally this makes writing the kernel easier, since we don’t have to worry about gaps or the ordering of the dims when indexing.
If we want to avoid this copy, `metal_kernel` automatically passes `a_shape`, `a_strides` and `a_ndim` for each input array `a` if any are present in `source`. We can then use MLX’s built in indexing utils to fetch the right elements for each thread.
Let’s convert `myexp` above to support arbitrarily strided arrays without relying on a copy from `ensure_row_contiguous`:
def exp\_elementwise(a: mx.array):
source \= """
uint elem = thread\_position\_in\_grid.x;
// Utils from \`mlx/backend/metal/kernels/utils.h\` are automatically included
uint loc = elem\_to\_loc(elem, inp\_shape, inp\_strides, inp\_ndim);
T tmp = inp\[loc\];
// Output arrays are always row contiguous
out\[elem\] = metal::exp(tmp);
"""
kernel \= mx.fast.metal\_kernel(
name\="myexp\_strided",
input\_names\=\["inp"\],
output\_names\=\["out"\],
source\=source
)
outputs \= kernel(
inputs\=\[a\],
template\=\[("T", mx.float32)\],
grid\=(a.size, 1, 1),
threadgroup\=(256, 1, 1),
output\_shapes\=\[a.shape\],
output\_dtypes\=\[a.dtype\],
ensure\_row\_contiguous\=False,
)
return outputs\[0\]
a \= mx.random.normal(shape\=(4, 16)).astype(mx.float16)
\# make non-contiguous
a \= a\[::2\]
b \= exp\_elementwise(a)
assert mx.allclose(b, mx.exp(a))
Complex Example[#](#complex-example "Link to this heading")
------------------------------------------------------------
Let’s implement a more complex example: `grid_sample` in `"bilinear"` mode.
We’ll start with the following MLX implementation using standard ops:
def grid\_sample\_ref(x, grid):
N, H\_in, W\_in, \_ \= x.shape
ix \= ((grid\[..., 0\] + 1) \* W\_in \- 1) / 2
iy \= ((grid\[..., 1\] + 1) \* H\_in \- 1) / 2
ix\_nw \= mx.floor(ix).astype(mx.int32)
iy\_nw \= mx.floor(iy).astype(mx.int32)
ix\_ne \= ix\_nw + 1
iy\_ne \= iy\_nw
ix\_sw \= ix\_nw
iy\_sw \= iy\_nw + 1
ix\_se \= ix\_nw + 1
iy\_se \= iy\_nw + 1
nw \= (ix\_se \- ix) \* (iy\_se \- iy)
ne \= (ix \- ix\_sw) \* (iy\_sw \- iy)
sw \= (ix\_ne \- ix) \* (iy \- iy\_ne)
se \= (ix \- ix\_nw) \* (iy \- iy\_nw)
I\_nw \= x\[mx.arange(N)\[:, None, None\], iy\_nw, ix\_nw, :\]
I\_ne \= x\[mx.arange(N)\[:, None, None\], iy\_ne, ix\_ne, :\]
I\_sw \= x\[mx.arange(N)\[:, None, None\], iy\_sw, ix\_sw, :\]
I\_se \= x\[mx.arange(N)\[:, None, None\], iy\_se, ix\_se, :\]
mask\_nw \= (iy\_nw \>= 0) & (iy\_nw <= H\_in \- 1) & (ix\_nw \>= 0) & (ix\_nw <= W\_in \- 1)
mask\_ne \= (iy\_ne \>= 0) & (iy\_ne <= H\_in \- 1) & (ix\_ne \>= 0) & (ix\_ne <= W\_in \- 1)
mask\_sw \= (iy\_sw \>= 0) & (iy\_sw <= H\_in \- 1) & (ix\_sw \>= 0) & (ix\_sw <= W\_in \- 1)
mask\_se \= (iy\_se \>= 0) & (iy\_se <= H\_in \- 1) & (ix\_se \>= 0) & (ix\_se <= W\_in \- 1)
I\_nw \*= mask\_nw\[..., None\]
I\_ne \*= mask\_ne\[..., None\]
I\_sw \*= mask\_sw\[..., None\]
I\_se \*= mask\_se\[..., None\]
output \= nw\[..., None\] \* I\_nw + ne\[..., None\] \* I\_ne + sw\[..., None\] \* I\_sw + se\[..., None\] \* I\_se
return output
Now let’s use `mx.custom_function` together with `mx.fast.metal_kernel` to write a fast GPU kernel for both the forward and backward passes.
First we’ll implement the forward pass as a fused kernel:
@mx.custom\_function
def grid\_sample(x, grid):
assert x.ndim \== 4, "\`x\` must be 4D."
assert grid.ndim \== 4, "\`grid\` must be 4D."
B, \_, \_, C \= x.shape
\_, gN, gM, D \= grid.shape
out\_shape \= (B, gN, gM, C)
assert D \== 2, "Last dim of \`grid\` must be size 2."
source \= """
uint elem = thread\_position\_in\_grid.x;
int H = x\_shape\[1\];
int W = x\_shape\[2\];
int C = x\_shape\[3\];
int gH = grid\_shape\[1\];
int gW = grid\_shape\[2\];
int w\_stride = C;
int h\_stride = W \* w\_stride;
int b\_stride = H \* h\_stride;
uint grid\_idx = elem / C \* 2;
float ix = ((grid\[grid\_idx\] + 1) \* W - 1) / 2;
float iy = ((grid\[grid\_idx + 1\] + 1) \* H - 1) / 2;
int ix\_nw = floor(ix);
int iy\_nw = floor(iy);
int ix\_ne = ix\_nw + 1;
int iy\_ne = iy\_nw;
int ix\_sw = ix\_nw;
int iy\_sw = iy\_nw + 1;
int ix\_se = ix\_nw + 1;
int iy\_se = iy\_nw + 1;
T nw = (ix\_se - ix) \* (iy\_se - iy);
T ne = (ix - ix\_sw) \* (iy\_sw - iy);
T sw = (ix\_ne - ix) \* (iy - iy\_ne);
T se = (ix - ix\_nw) \* (iy - iy\_nw);
int batch\_idx = elem / C / gH / gW \* b\_stride;
int channel\_idx = elem % C;
int base\_idx = batch\_idx + channel\_idx;
T I\_nw = x\[base\_idx + iy\_nw \* h\_stride + ix\_nw \* w\_stride\];
T I\_ne = x\[base\_idx + iy\_ne \* h\_stride + ix\_ne \* w\_stride\];
T I\_sw = x\[base\_idx + iy\_sw \* h\_stride + ix\_sw \* w\_stride\];
T I\_se = x\[base\_idx + iy\_se \* h\_stride + ix\_se \* w\_stride\];
I\_nw = iy\_nw >= 0 && iy\_nw <= H - 1 && ix\_nw >= 0 && ix\_nw <= W - 1 ? I\_nw : 0;
I\_ne = iy\_ne >= 0 && iy\_ne <= H - 1 && ix\_ne >= 0 && ix\_ne <= W - 1 ? I\_ne : 0;
I\_sw = iy\_sw >= 0 && iy\_sw <= H - 1 && ix\_sw >= 0 && ix\_sw <= W - 1 ? I\_sw : 0;
I\_se = iy\_se >= 0 && iy\_se <= H - 1 && ix\_se >= 0 && ix\_se <= W - 1 ? I\_se : 0;
out\[elem\] = nw \* I\_nw + ne \* I\_ne + sw \* I\_sw + se \* I\_se;
"""
kernel \= mx.fast.metal\_kernel(
name\="grid\_sample",
input\_names\=\["x", "grid"\],
output\_names\=\["out"\],
source\=source,
)
outputs \= kernel(
inputs\=\[x, grid\],
template\=\[("T", x.dtype)\],
output\_shapes\=\[out\_shape\],
output\_dtypes\=\[x.dtype\],
grid\=(np.prod(out\_shape), 1, 1),
threadgroup\=(256, 1, 1),
)
return outputs\[0\]
For a reasonably sized input such as:
x.shape \= (8, 1024, 1024, 64)
grid.shape \= (8, 256, 256, 2)
On an M1 Max, we see a big performance improvement:
`55.7ms -> 6.7ms => 8x speed up`
Grid Sample VJP[#](#grid-sample-vjp "Link to this heading")
------------------------------------------------------------
Since we decorated `grid_sample` with `mx.custom_function`, we can now define its custom vjp transform so MLX can differentiate it.
The backwards pass requires atomically updating `x_grad`/`grid_grad` and so requires a few extra `mx.fast.metal_kernel` features:
* `init_value=0`
Initialize all of the kernel’s outputs to this value before it runs. This allows us to update only part of the output arrays with the kernel.
* `atomic_outputs=True`
Designate all of the kernel outputs as `atomic` in the function signature. This means we can use Metal’s `atomic` features to simultaneously update the `x_grad` and `grid_grad` arrays from multiple threadgroups. See section 6.15 of the [Metal Shading Language Specification](https://developer.apple.com/metal/Metal-Shading-Language-Specification.pdf)
for more details.
We can then implement the backwards pass as follows:
@grid\_sample.vjp
def grid\_sample\_vjp(primals, cotangent, \_):
x, grid \= primals
B, \_, \_, C \= x.shape
\_, gN, gM, D \= grid.shape
assert D \== 2, "Last dim of \`grid\` must be size 2."
source \= """
uint elem = thread\_position\_in\_grid.x;
int H = x\_shape\[1\];
int W = x\_shape\[2\];
int C = x\_shape\[3\];
// Pad C to the nearest larger simdgroup size multiple
int C\_padded = ceildiv(C, threads\_per\_simdgroup) \* threads\_per\_simdgroup;
int gH = grid\_shape\[1\];
int gW = grid\_shape\[2\];
int w\_stride = C;
int h\_stride = W \* w\_stride;
int b\_stride = H \* h\_stride;
uint grid\_idx = elem / C\_padded \* 2;
float ix = ((grid\[grid\_idx\] + 1) \* W - 1) / 2;
float iy = ((grid\[grid\_idx + 1\] + 1) \* H - 1) / 2;
int ix\_nw = floor(ix);
int iy\_nw = floor(iy);
int ix\_ne = ix\_nw + 1;
int iy\_ne = iy\_nw;
int ix\_sw = ix\_nw;
int iy\_sw = iy\_nw + 1;
int ix\_se = ix\_nw + 1;
int iy\_se = iy\_nw + 1;
T nw = (ix\_se - ix) \* (iy\_se - iy);
T ne = (ix - ix\_sw) \* (iy\_sw - iy);
T sw = (ix\_ne - ix) \* (iy - iy\_ne);
T se = (ix - ix\_nw) \* (iy - iy\_nw);
int batch\_idx = elem / C\_padded / gH / gW \* b\_stride;
int channel\_idx = elem % C\_padded;
int base\_idx = batch\_idx + channel\_idx;
T gix = T(0);
T giy = T(0);
if (channel\_idx < C) {
int cot\_index = elem / C\_padded \* C + channel\_idx;
T cot = cotangent\[cot\_index\];
if (iy\_nw >= 0 && iy\_nw <= H - 1 && ix\_nw >= 0 && ix\_nw <= W - 1) {
int offset = base\_idx + iy\_nw \* h\_stride + ix\_nw \* w\_stride;
atomic\_fetch\_add\_explicit(&x\_grad\[offset\], nw \* cot, memory\_order\_relaxed);
T I\_nw = x\[offset\];
gix -= I\_nw \* (iy\_se - iy) \* cot;
giy -= I\_nw \* (ix\_se - ix) \* cot;
}
if (iy\_ne >= 0 && iy\_ne <= H - 1 && ix\_ne >= 0 && ix\_ne <= W - 1) {
int offset = base\_idx + iy\_ne \* h\_stride + ix\_ne \* w\_stride;
atomic\_fetch\_add\_explicit(&x\_grad\[offset\], ne \* cot, memory\_order\_relaxed);
T I\_ne = x\[offset\];
gix += I\_ne \* (iy\_sw - iy) \* cot;
giy -= I\_ne \* (ix - ix\_sw) \* cot;
}
if (iy\_sw >= 0 && iy\_sw <= H - 1 && ix\_sw >= 0 && ix\_sw <= W - 1) {
int offset = base\_idx + iy\_sw \* h\_stride + ix\_sw \* w\_stride;
atomic\_fetch\_add\_explicit(&x\_grad\[offset\], sw \* cot, memory\_order\_relaxed);
T I\_sw = x\[offset\];
gix -= I\_sw \* (iy - iy\_ne) \* cot;
giy += I\_sw \* (ix\_ne - ix) \* cot;
}
if (iy\_se >= 0 && iy\_se <= H - 1 && ix\_se >= 0 && ix\_se <= W - 1) {
int offset = base\_idx + iy\_se \* h\_stride + ix\_se \* w\_stride;
atomic\_fetch\_add\_explicit(&x\_grad\[offset\], se \* cot, memory\_order\_relaxed);
T I\_se = x\[offset\];
gix += I\_se \* (iy - iy\_nw) \* cot;
giy += I\_se \* (ix - ix\_nw) \* cot;
}
}
T gix\_mult = W / 2;
T giy\_mult = H / 2;
// Reduce across each simdgroup first.
// This is much faster than relying purely on atomics.
gix = simd\_sum(gix);
giy = simd\_sum(giy);
if (thread\_index\_in\_simdgroup == 0) {
atomic\_fetch\_add\_explicit(&grid\_grad\[grid\_idx\], gix \* gix\_mult, memory\_order\_relaxed);
atomic\_fetch\_add\_explicit(&grid\_grad\[grid\_idx + 1\], giy \* giy\_mult, memory\_order\_relaxed);
}
"""
kernel \= mx.fast.metal\_kernel(
name\="grid\_sample\_grad",
input\_names\=\["x", "grid", "cotangent"\],
output\_names\=\["x\_grad", "grid\_grad"\],
source\=source,
atomic\_outputs\=True,
)
\# pad the output channels to simd group size
\# so that our \`simd\_sum\`s don't overlap.
simdgroup\_size \= 32
C\_padded \= (C + simdgroup\_size \- 1) // simdgroup\_size \* simdgroup\_size
grid\_size \= B \* gN \* gM \* C\_padded
outputs \= kernel(
inputs\=\[x, grid, cotangent\],
template\=\[("T", x.dtype)\],
output\_shapes\=\[x.shape, grid.shape\],
output\_dtypes\=\[x.dtype, x.dtype\],
grid\=(grid\_size, 1, 1),
threadgroup\=(256, 1, 1),
init\_value\=0,
)
return outputs\[0\], outputs\[1\]
There’s an even larger speed up for the vjp:
`676.4ms -> 16.7ms => 40x speed up`
Contents
---
# Operations — MLX 0.24.0 documentation
[Skip to main content](#main-content)
Back to top
Ctrl+K
[](../index.html)
[](https://github.com/ml-explore/mlx "Source repository")
* [.rst](../_sources/cpp/ops.rst "Download source file")
* .pdf
Operations
==========
Contents
--------
Operations[#](#operations "Link to this heading")
==================================================
array arange(double start, double stop, double step, Dtype dtype, StreamOrDevice s \= {})[#](#_CPPv46arangeddd5Dtype14StreamOrDevice "Link to this definition")
A 1D array of numbers starting at `start` (optional), stopping at stop, stepping by `step` (optional).
array arange(double start, double stop, double step, StreamOrDevice s \= {})[#](#_CPPv46arangeddd14StreamOrDevice "Link to this definition")
array arange(double start, double stop, Dtype dtype, StreamOrDevice s \= {})[#](#_CPPv46arangedd5Dtype14StreamOrDevice "Link to this definition")
array arange(double start, double stop, StreamOrDevice s \= {})[#](#_CPPv46arangedd14StreamOrDevice "Link to this definition")
array arange(double stop, Dtype dtype, StreamOrDevice s \= {})[#](#_CPPv46aranged5Dtype14StreamOrDevice "Link to this definition")
array arange(double stop, StreamOrDevice s \= {})[#](#_CPPv46aranged14StreamOrDevice "Link to this definition")
array arange(int start, int stop, int step, StreamOrDevice s \= {})[#](#_CPPv46arangeiii14StreamOrDevice "Link to this definition")
array arange(int start, int stop, StreamOrDevice s \= {})[#](#_CPPv46arangeii14StreamOrDevice "Link to this definition")
array arange(int stop, StreamOrDevice s \= {})[#](#_CPPv46arangei14StreamOrDevice "Link to this definition")
array linspace(double start, double stop, int num \= 50, Dtype dtype \= float32, StreamOrDevice s \= {})[#](#_CPPv48linspaceddi5Dtype14StreamOrDevice "Link to this definition")
A 1D array of `num` evenly spaced numbers in the range `[start, stop]`
array astype(array a, Dtype dtype, StreamOrDevice s \= {})[#](#_CPPv46astype5array5Dtype14StreamOrDevice "Link to this definition")
Convert an array to the given data type.
array as\_strided(array a, Shape shape, Strides strides, size\_t offset, StreamOrDevice s \= {})[#](#_CPPv410as_strided5array5Shape7Strides6size_t14StreamOrDevice "Link to this definition")
Create a view of an array with the given shape and strides.
array copy(array a, StreamOrDevice s \= {})[#](#_CPPv44copy5array14StreamOrDevice "Link to this definition")
Copy another array.
array full(Shape shape, array vals, Dtype dtype, StreamOrDevice s \= {})[#](#_CPPv44full5Shape5array5Dtype14StreamOrDevice "Link to this definition")
Fill an array of the given shape with the given value(s).
array full(Shape shape, array vals, StreamOrDevice s \= {})[#](#_CPPv44full5Shape5array14StreamOrDevice "Link to this definition")
template
array full(Shape shape, [T](#_CPPv4I0E4full5array5Shape1T5Dtype14StreamOrDevice "full::T")
val, Dtype dtype, StreamOrDevice s \= {})[#](#_CPPv4I0E4full5array5Shape1T5Dtype14StreamOrDevice "Link to this definition")
template
array full(Shape shape, [T](#_CPPv4I0E4full5array5Shape1T14StreamOrDevice "full::T")
val, StreamOrDevice s \= {})[#](#_CPPv4I0E4full5array5Shape1T14StreamOrDevice "Link to this definition")
array zeros(const Shape &shape, Dtype dtype, StreamOrDevice s \= {})[#](#_CPPv45zerosRK5Shape5Dtype14StreamOrDevice "Link to this definition")
Fill an array of the given shape with zeros.
inline array zeros(const Shape &shape, StreamOrDevice s \= {})[#](#_CPPv45zerosRK5Shape14StreamOrDevice "Link to this definition")
array zeros\_like(const array &a, StreamOrDevice s \= {})[#](#_CPPv410zeros_likeRK5array14StreamOrDevice "Link to this definition")
array ones(const Shape &shape, Dtype dtype, StreamOrDevice s \= {})[#](#_CPPv44onesRK5Shape5Dtype14StreamOrDevice "Link to this definition")
Fill an array of the given shape with ones.
inline array ones(const Shape &shape, StreamOrDevice s \= {})[#](#_CPPv44onesRK5Shape14StreamOrDevice "Link to this definition")
array ones\_like(const array &a, StreamOrDevice s \= {})[#](#_CPPv49ones_likeRK5array14StreamOrDevice "Link to this definition")
array eye(int n, int m, int k, Dtype dtype, StreamOrDevice s \= {})[#](#_CPPv43eyeiii5Dtype14StreamOrDevice "Link to this definition")
Fill an array of the given shape (n,m) with ones in the specified diagonal k, and zeros everywhere else.
inline array eye(int n, Dtype dtype, StreamOrDevice s \= {})[#](#_CPPv43eyei5Dtype14StreamOrDevice "Link to this definition")
inline array eye(int n, int m, StreamOrDevice s \= {})[#](#_CPPv43eyeii14StreamOrDevice "Link to this definition")
inline array eye(int n, int m, int k, StreamOrDevice s \= {})[#](#_CPPv43eyeiii14StreamOrDevice "Link to this definition")
inline array eye(int n, StreamOrDevice s \= {})[#](#_CPPv43eyei14StreamOrDevice "Link to this definition")
array identity(int n, Dtype dtype, StreamOrDevice s \= {})[#](#_CPPv48identityi5Dtype14StreamOrDevice "Link to this definition")
Create a square matrix of shape (n,n) of zeros, and ones in the major diagonal.
inline array identity(int n, StreamOrDevice s \= {})[#](#_CPPv48identityi14StreamOrDevice "Link to this definition")
array tri(int n, int m, int k, Dtype type, StreamOrDevice s \= {})[#](#_CPPv43triiii5Dtype14StreamOrDevice "Link to this definition")
inline array tri(int n, Dtype type, StreamOrDevice s \= {})[#](#_CPPv43trii5Dtype14StreamOrDevice "Link to this definition")
array tril(array x, int k \= 0, StreamOrDevice s \= {})[#](#_CPPv44tril5arrayi14StreamOrDevice "Link to this definition")
array triu(array x, int k \= 0, StreamOrDevice s \= {})[#](#_CPPv44triu5arrayi14StreamOrDevice "Link to this definition")
array reshape(const array &a, Shape shape, StreamOrDevice s \= {})[#](#_CPPv47reshapeRK5array5Shape14StreamOrDevice "Link to this definition")
Reshape an array to the given shape.
array unflatten(const array &a, int axis, Shape shape, StreamOrDevice s \= {})[#](#_CPPv49unflattenRK5arrayi5Shape14StreamOrDevice "Link to this definition")
Unflatten the axis to the given shape.
array flatten(const array &a, int start\_axis, int end\_axis \= \-1, StreamOrDevice s \= {})[#](#_CPPv47flattenRK5arrayii14StreamOrDevice "Link to this definition")
Flatten the dimensions in the range `[start_axis, end_axis]` .
array flatten(const array &a, StreamOrDevice s \= {})[#](#_CPPv47flattenRK5array14StreamOrDevice "Link to this definition")
Flatten the array to 1D.
array hadamard\_transform(const array &a, [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::optional scale \= [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::nullopt, StreamOrDevice s \= {})[#](#_CPPv418hadamard_transformRK5arrayNSt8optionalIfEE14StreamOrDevice "Link to this definition")
Multiply the array by the Hadamard matrix of corresponding size.
array squeeze(const array &a, const [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector &axes, StreamOrDevice s \= {})[#](#_CPPv47squeezeRK5arrayRKNSt6vectorIiEE14StreamOrDevice "Link to this definition")
Remove singleton dimensions at the given axes.
array squeeze(const array &a, int axis, StreamOrDevice s \= {})[#](#_CPPv47squeezeRK5arrayi14StreamOrDevice "Link to this definition")
Remove singleton dimensions at the given axis.
array squeeze(const array &a, StreamOrDevice s \= {})[#](#_CPPv47squeezeRK5array14StreamOrDevice "Link to this definition")
Remove all singleton dimensions.
array expand\_dims(const array &a, const [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector &axes, StreamOrDevice s \= {})[#](#_CPPv411expand_dimsRK5arrayRKNSt6vectorIiEE14StreamOrDevice "Link to this definition")
Add a singleton dimension at the given axes.
array expand\_dims(const array &a, int axis, StreamOrDevice s \= {})[#](#_CPPv411expand_dimsRK5arrayi14StreamOrDevice "Link to this definition")
Add a singleton dimension at the given axis.
array slice(const array &a, Shape start, Shape stop, Shape strides, StreamOrDevice s \= {})[#](#_CPPv45sliceRK5array5Shape5Shape5Shape14StreamOrDevice "Link to this definition")
Slice an array.
inline array slice(const array &a, [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::initializer\_list start, Shape stop, Shape strides, StreamOrDevice s \= {})[#](#_CPPv45sliceRK5arrayNSt16initializer_listIiEE5Shape5Shape14StreamOrDevice "Link to this definition")
array slice(const array &a, Shape start, Shape stop, StreamOrDevice s \= {})[#](#_CPPv45sliceRK5array5Shape5Shape14StreamOrDevice "Link to this definition")
Slice an array with a stride of 1 in each dimension.
array slice(const array &a, const array &start, [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector axes, Shape slice\_size, StreamOrDevice s \= {})[#](#_CPPv45sliceRK5arrayRK5arrayNSt6vectorIiEE5Shape14StreamOrDevice "Link to this definition")
Slice an array with dynamic starting indices.
array slice\_update(const array &src, const array &update, Shape start, Shape stop, Shape strides, StreamOrDevice s \= {})[#](#_CPPv412slice_updateRK5arrayRK5array5Shape5Shape5Shape14StreamOrDevice "Link to this definition")
Update a slice from the source array.
array slice\_update(const array &src, const array &update, Shape start, Shape stop, StreamOrDevice s \= {})[#](#_CPPv412slice_updateRK5arrayRK5array5Shape5Shape14StreamOrDevice "Link to this definition")
Update a slice from the source array with stride 1 in each dimension.
array slice\_update(const array &src, const array &update, const array &start, [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector axes, StreamOrDevice s \= {})[#](#_CPPv412slice_updateRK5arrayRK5arrayRK5arrayNSt6vectorIiEE14StreamOrDevice "Link to this definition")
Update a slice from the source array with dynamic starting indices.
[std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector split(const array &a, int num\_splits, int axis, StreamOrDevice s \= {})[#](#_CPPv45splitRK5arrayii14StreamOrDevice "Link to this definition")
Split an array into sub-arrays along a given axis.
[std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector split(const array &a, int num\_splits, StreamOrDevice s \= {})[#](#_CPPv45splitRK5arrayi14StreamOrDevice "Link to this definition")
[std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector split(const array &a, const Shape &indices, int axis, StreamOrDevice s \= {})[#](#_CPPv45splitRK5arrayRK5Shapei14StreamOrDevice "Link to this definition")
[std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector split(const array &a, const Shape &indices, StreamOrDevice s \= {})[#](#_CPPv45splitRK5arrayRK5Shape14StreamOrDevice "Link to this definition")
[std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector meshgrid(const [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector &arrays, bool sparse \= false, const [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::string &indexing \= "xy", StreamOrDevice s \= {})[#](#_CPPv48meshgridRKNSt6vectorI5arrayEEbRKNSt6stringE14StreamOrDevice "Link to this definition")
A vector of coordinate arrays from coordinate vectors.
array clip(const array &a, const [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::optional &a\_min \= [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::nullopt, const [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::optional &a\_max \= [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::nullopt, StreamOrDevice s \= {})[#](#_CPPv44clipRK5arrayRKNSt8optionalI5arrayEERKNSt8optionalI5arrayEE14StreamOrDevice "Link to this definition")
Clip (limit) the values in an array.
array concatenate([std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector arrays, int axis, StreamOrDevice s \= {})[#](#_CPPv411concatenateNSt6vectorI5arrayEEi14StreamOrDevice "Link to this definition")
Concatenate arrays along a given axis.
array concatenate([std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector arrays, StreamOrDevice s \= {})[#](#_CPPv411concatenateNSt6vectorI5arrayEE14StreamOrDevice "Link to this definition")
array stack(const [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector &arrays, int axis, StreamOrDevice s \= {})[#](#_CPPv45stackRKNSt6vectorI5arrayEEi14StreamOrDevice "Link to this definition")
Stack arrays along a new axis.
array stack(const [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector &arrays, StreamOrDevice s \= {})[#](#_CPPv45stackRKNSt6vectorI5arrayEE14StreamOrDevice "Link to this definition")
array repeat(const array &arr, int repeats, int axis, StreamOrDevice s \= {})[#](#_CPPv46repeatRK5arrayii14StreamOrDevice "Link to this definition")
Repeat an array along an axis.
array repeat(const array &arr, int repeats, StreamOrDevice s \= {})[#](#_CPPv46repeatRK5arrayi14StreamOrDevice "Link to this definition")
array tile(const array &arr, [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector reps, StreamOrDevice s \= {})[#](#_CPPv44tileRK5arrayNSt6vectorIiEE14StreamOrDevice "Link to this definition")
array transpose(const array &a, [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector axes, StreamOrDevice s \= {})[#](#_CPPv49transposeRK5arrayNSt6vectorIiEE14StreamOrDevice "Link to this definition")
Permutes the dimensions according to the given axes.
inline array transpose(const array &a, [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::initializer\_list axes, StreamOrDevice s \= {})[#](#_CPPv49transposeRK5arrayNSt16initializer_listIiEE14StreamOrDevice "Link to this definition")
array swapaxes(const array &a, int axis1, int axis2, StreamOrDevice s \= {})[#](#_CPPv48swapaxesRK5arrayii14StreamOrDevice "Link to this definition")
Swap two axes of an array.
array moveaxis(const array &a, int source, int destination, StreamOrDevice s \= {})[#](#_CPPv48moveaxisRK5arrayii14StreamOrDevice "Link to this definition")
Move an axis of an array.
array pad(const array &a, const [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector &axes, const Shape &low\_pad\_size, const Shape &high\_pad\_size, const array &pad\_value \= array(0), const [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::string &mode \= "constant", StreamOrDevice s \= {})[#](#_CPPv43padRK5arrayRKNSt6vectorIiEERK5ShapeRK5ShapeRK5arrayRKNSt6stringE14StreamOrDevice "Link to this definition")
Pad an array with a constant value.
array pad(const array &a, const [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector<[std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::pair\> &pad\_width, const array &pad\_value \= array(0), const [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::string &mode \= "constant", StreamOrDevice s \= {})[#](#_CPPv43padRK5arrayRKNSt6vectorINSt4pairIiiEEEERK5arrayRKNSt6stringE14StreamOrDevice "Link to this definition")
Pad an array with a constant value along all axes.
array pad(const array &a, const [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::pair &pad\_width, const array &pad\_value \= array(0), const [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::string &mode \= "constant", StreamOrDevice s \= {})[#](#_CPPv43padRK5arrayRKNSt4pairIiiEERK5arrayRKNSt6stringE14StreamOrDevice "Link to this definition")
array pad(const array &a, int pad\_width, const array &pad\_value \= array(0), const [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::string &mode \= "constant", StreamOrDevice s \= {})[#](#_CPPv43padRK5arrayiRK5arrayRKNSt6stringE14StreamOrDevice "Link to this definition")
array transpose(const array &a, StreamOrDevice s \= {})[#](#_CPPv49transposeRK5array14StreamOrDevice "Link to this definition")
Permutes the dimensions in reverse order.
array broadcast\_to(const array &a, const Shape &shape, StreamOrDevice s \= {})[#](#_CPPv412broadcast_toRK5arrayRK5Shape14StreamOrDevice "Link to this definition")
Broadcast an array to a given shape.
[std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector broadcast\_arrays(const [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector &inputs, StreamOrDevice s \= {})[#](#_CPPv416broadcast_arraysRKNSt6vectorI5arrayEE14StreamOrDevice "Link to this definition")
Broadcast a vector of arrays against one another.
array equal(const array &a, const array &b, StreamOrDevice s \= {})[#](#_CPPv45equalRK5arrayRK5array14StreamOrDevice "Link to this definition")
Returns the bool array with (a == b) element-wise.
inline array operator\==(const array &a, const array &b)[#](#_CPPv4eqRK5arrayRK5array "Link to this definition")
template
array operator\==([T](#_CPPv4I0Eeq5array1TRK5array "operator==::T")
a, const array &b)[#](#_CPPv4I0Eeq5array1TRK5array "Link to this definition")
template
array operator\==(const array &a, [T](#_CPPv4I0Eeq5arrayRK5array1T "operator==::T")
b)[#](#_CPPv4I0Eeq5arrayRK5array1T "Link to this definition")
array not\_equal(const array &a, const array &b, StreamOrDevice s \= {})[#](#_CPPv49not_equalRK5arrayRK5array14StreamOrDevice "Link to this definition")
Returns the bool array with (a != b) element-wise.
inline array operator!=(const array &a, const array &b)[#](#_CPPv4neRK5arrayRK5array "Link to this definition")
template
array operator!=([T](#_CPPv4I0Ene5array1TRK5array "operator!=::T")
a, const array &b)[#](#_CPPv4I0Ene5array1TRK5array "Link to this definition")
template
array operator!=(const array &a, [T](#_CPPv4I0Ene5arrayRK5array1T "operator!=::T")
b)[#](#_CPPv4I0Ene5arrayRK5array1T "Link to this definition")
array greater(const array &a, const array &b, StreamOrDevice s \= {})[#](#_CPPv47greaterRK5arrayRK5array14StreamOrDevice "Link to this definition")
Returns bool array with (a > b) element-wise.
inline array operator\>(const array &a, const array &b)[#](#_CPPv4gtRK5arrayRK5array "Link to this definition")
template
array operator\>([T](#_CPPv4I0Egt5array1TRK5array "operator>::T")
a, const array &b)[#](#_CPPv4I0Egt5array1TRK5array "Link to this definition")
template
array operator\>(const array &a, [T](#_CPPv4I0Egt5arrayRK5array1T "operator>::T")
b)[#](#_CPPv4I0Egt5arrayRK5array1T "Link to this definition")
array greater\_equal(const array &a, const array &b, StreamOrDevice s \= {})[#](#_CPPv413greater_equalRK5arrayRK5array14StreamOrDevice "Link to this definition")
Returns bool array with (a >= b) element-wise.
inline array operator\>=(const array &a, const array &b)[#](#_CPPv4geRK5arrayRK5array "Link to this definition")
template
array operator\>=([T](#_CPPv4I0Ege5array1TRK5array "operator>=::T")
a, const array &b)[#](#_CPPv4I0Ege5array1TRK5array "Link to this definition")
template
array operator\>=(const array &a, [T](#_CPPv4I0Ege5arrayRK5array1T "operator>=::T")
b)[#](#_CPPv4I0Ege5arrayRK5array1T "Link to this definition")
array less(const array &a, const array &b, StreamOrDevice s \= {})[#](#_CPPv44lessRK5arrayRK5array14StreamOrDevice "Link to this definition")
Returns bool array with (a < b) element-wise.
inline array operator<(const array &a, const array &b)[#](#_CPPv4ltRK5arrayRK5array "Link to this definition")
template
array operator<([T](#_CPPv4I0Elt5array1TRK5array "operator<::T")
a, const array &b)[#](#_CPPv4I0Elt5array1TRK5array "Link to this definition")
template
array operator<(const array &a, [T](#_CPPv4I0Elt5arrayRK5array1T "operator<::T")
b)[#](#_CPPv4I0Elt5arrayRK5array1T "Link to this definition")
array less\_equal(const array &a, const array &b, StreamOrDevice s \= {})[#](#_CPPv410less_equalRK5arrayRK5array14StreamOrDevice "Link to this definition")
Returns bool array with (a <= b) element-wise.
inline array operator<=(const array &a, const array &b)[#](#_CPPv4leRK5arrayRK5array "Link to this definition")
template
array operator<=([T](#_CPPv4I0Ele5array1TRK5array "operator<=::T")
a, const array &b)[#](#_CPPv4I0Ele5array1TRK5array "Link to this definition")
template
array operator<=(const array &a, [T](#_CPPv4I0Ele5arrayRK5array1T "operator<=::T")
b)[#](#_CPPv4I0Ele5arrayRK5array1T "Link to this definition")
array array\_equal(const array &a, const array &b, bool equal\_nan, StreamOrDevice s \= {})[#](#_CPPv411array_equalRK5arrayRK5arrayb14StreamOrDevice "Link to this definition")
True if two arrays have the same shape and elements.
inline array array\_equal(const array &a, const array &b, StreamOrDevice s \= {})[#](#_CPPv411array_equalRK5arrayRK5array14StreamOrDevice "Link to this definition")
array isnan(const array &a, StreamOrDevice s \= {})[#](#_CPPv45isnanRK5array14StreamOrDevice "Link to this definition")
array isinf(const array &a, StreamOrDevice s \= {})[#](#_CPPv45isinfRK5array14StreamOrDevice "Link to this definition")
array isfinite(const array &a, StreamOrDevice s \= {})[#](#_CPPv48isfiniteRK5array14StreamOrDevice "Link to this definition")
array isposinf(const array &a, StreamOrDevice s \= {})[#](#_CPPv48isposinfRK5array14StreamOrDevice "Link to this definition")
array isneginf(const array &a, StreamOrDevice s \= {})[#](#_CPPv48isneginfRK5array14StreamOrDevice "Link to this definition")
array where(const array &condition, const array &x, const array &y, StreamOrDevice s \= {})[#](#_CPPv45whereRK5arrayRK5arrayRK5array14StreamOrDevice "Link to this definition")
Select from x or y depending on condition.
array nan\_to\_num(const array &a, float nan \= 0.0f, const [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::optional posinf \= [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::nullopt, const [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::optional neginf \= [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::nullopt, StreamOrDevice s \= {})[#](#_CPPv410nan_to_numRK5arrayfKNSt8optionalIfEEKNSt8optionalIfEE14StreamOrDevice "Link to this definition")
Replace NaN and infinities with finite numbers.
array all(const array &a, bool keepdims, StreamOrDevice s \= {})[#](#_CPPv43allRK5arrayb14StreamOrDevice "Link to this definition")
True if all elements in the array are true (or non-zero).
inline array all(const array &a, StreamOrDevice s \= {})[#](#_CPPv43allRK5array14StreamOrDevice "Link to this definition")
array allclose(const array &a, const array &b, double rtol \= 1e-5, double atol \= 1e-8, bool equal\_nan \= false, StreamOrDevice s \= {})[#](#_CPPv48allcloseRK5arrayRK5arrayddb14StreamOrDevice "Link to this definition")
True if the two arrays are equal within the specified tolerance.
array isclose(const array &a, const array &b, double rtol \= 1e-5, double atol \= 1e-8, bool equal\_nan \= false, StreamOrDevice s \= {})[#](#_CPPv47iscloseRK5arrayRK5arrayddb14StreamOrDevice "Link to this definition")
Returns a boolean array where two arrays are element-wise equal within the specified tolerance.
array all(const array &a, const [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector &axes, bool keepdims \= false, StreamOrDevice s \= {})[#](#_CPPv43allRK5arrayRKNSt6vectorIiEEb14StreamOrDevice "Link to this definition")
Reduces the input along the given axes.
An output value is true if all the corresponding inputs are true.
array all(const array &a, int axis, bool keepdims \= false, StreamOrDevice s \= {})[#](#_CPPv43allRK5arrayib14StreamOrDevice "Link to this definition")
Reduces the input along the given axis.
An output value is true if all the corresponding inputs are true.
array any(const array &a, bool keepdims, StreamOrDevice s \= {})[#](#_CPPv43anyRK5arrayb14StreamOrDevice "Link to this definition")
True if any elements in the array are true (or non-zero).
inline array any(const array &a, StreamOrDevice s \= {})[#](#_CPPv43anyRK5array14StreamOrDevice "Link to this definition")
array any(const array &a, const [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector &axes, bool keepdims \= false, StreamOrDevice s \= {})[#](#_CPPv43anyRK5arrayRKNSt6vectorIiEEb14StreamOrDevice "Link to this definition")
Reduces the input along the given axes.
An output value is true if any of the corresponding inputs are true.
array any(const array &a, int axis, bool keepdims \= false, StreamOrDevice s \= {})[#](#_CPPv43anyRK5arrayib14StreamOrDevice "Link to this definition")
Reduces the input along the given axis.
An output value is true if any of the corresponding inputs are true.
array sum(const array &a, bool keepdims, StreamOrDevice s \= {})[#](#_CPPv43sumRK5arrayb14StreamOrDevice "Link to this definition")
Sums the elements of an array.
inline array sum(const array &a, StreamOrDevice s \= {})[#](#_CPPv43sumRK5array14StreamOrDevice "Link to this definition")
array sum(const array &a, const [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector &axes, bool keepdims \= false, StreamOrDevice s \= {})[#](#_CPPv43sumRK5arrayRKNSt6vectorIiEEb14StreamOrDevice "Link to this definition")
Sums the elements of an array along the given axes.
array sum(const array &a, int axis, bool keepdims \= false, StreamOrDevice s \= {})[#](#_CPPv43sumRK5arrayib14StreamOrDevice "Link to this definition")
Sums the elements of an array along the given axis.
array mean(const array &a, bool keepdims, StreamOrDevice s \= {})[#](#_CPPv44meanRK5arrayb14StreamOrDevice "Link to this definition")
Computes the mean of the elements of an array.
inline array mean(const array &a, StreamOrDevice s \= {})[#](#_CPPv44meanRK5array14StreamOrDevice "Link to this definition")
array mean(const array &a, const [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector &axes, bool keepdims \= false, StreamOrDevice s \= {})[#](#_CPPv44meanRK5arrayRKNSt6vectorIiEEb14StreamOrDevice "Link to this definition")
Computes the mean of the elements of an array along the given axes.
array mean(const array &a, int axis, bool keepdims \= false, StreamOrDevice s \= {})[#](#_CPPv44meanRK5arrayib14StreamOrDevice "Link to this definition")
Computes the mean of the elements of an array along the given axis.
array var(const array &a, bool keepdims, int ddof \= 0, StreamOrDevice s \= {})[#](#_CPPv43varRK5arraybi14StreamOrDevice "Link to this definition")
Computes the variance of the elements of an array.
inline array var(const array &a, StreamOrDevice s \= {})[#](#_CPPv43varRK5array14StreamOrDevice "Link to this definition")
array var(const array &a, const [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector &axes, bool keepdims \= false, int ddof \= 0, StreamOrDevice s \= {})[#](#_CPPv43varRK5arrayRKNSt6vectorIiEEbi14StreamOrDevice "Link to this definition")
Computes the variance of the elements of an array along the given axes.
array var(const array &a, int axis, bool keepdims \= false, int ddof \= 0, StreamOrDevice s \= {})[#](#_CPPv43varRK5arrayibi14StreamOrDevice "Link to this definition")
Computes the variance of the elements of an array along the given axis.
array std(const array &a, bool keepdims, int ddof \= 0, StreamOrDevice s \= {})[#](#_CPPv4StRK5arraybi14StreamOrDevice "Link to this definition")
Computes the standard deviation of the elements of an array.
inline array std(const array &a, StreamOrDevice s \= {})[#](#_CPPv4StRK5array14StreamOrDevice "Link to this definition")
array std(const array &a, const [std](#_CPPv4StRK5arrayRKNSt6vectorIiEEbi14StreamOrDevice "std")
::vector &axes, bool keepdims \= false, int ddof \= 0, StreamOrDevice s \= {})[#](#_CPPv4StRK5arrayRKNSt6vectorIiEEbi14StreamOrDevice "Link to this definition")
Computes the standard deviatoin of the elements of an array along the given axes.
array std(const array &a, int axis, bool keepdims \= false, int ddof \= 0, StreamOrDevice s \= {})[#](#_CPPv4StRK5arrayibi14StreamOrDevice "Link to this definition")
Computes the standard deviation of the elements of an array along the given axis.
array prod(const array &a, bool keepdims, StreamOrDevice s \= {})[#](#_CPPv44prodRK5arrayb14StreamOrDevice "Link to this definition")
The product of all elements of the array.
inline array prod(const array &a, StreamOrDevice s \= {})[#](#_CPPv44prodRK5array14StreamOrDevice "Link to this definition")
array prod(const array &a, const [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector &axes, bool keepdims \= false, StreamOrDevice s \= {})[#](#_CPPv44prodRK5arrayRKNSt6vectorIiEEb14StreamOrDevice "Link to this definition")
The product of the elements of an array along the given axes.
array prod(const array &a, int axis, bool keepdims \= false, StreamOrDevice s \= {})[#](#_CPPv44prodRK5arrayib14StreamOrDevice "Link to this definition")
The product of the elements of an array along the given axis.
array max(const array &a, bool keepdims, StreamOrDevice s \= {})[#](#_CPPv43maxRK5arrayb14StreamOrDevice "Link to this definition")
The maximum of all elements of the array.
inline array max(const array &a, StreamOrDevice s \= {})[#](#_CPPv43maxRK5array14StreamOrDevice "Link to this definition")
array max(const array &a, const [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector &axes, bool keepdims \= false, StreamOrDevice s \= {})[#](#_CPPv43maxRK5arrayRKNSt6vectorIiEEb14StreamOrDevice "Link to this definition")
The maximum of the elements of an array along the given axes.
array max(const array &a, int axis, bool keepdims \= false, StreamOrDevice s \= {})[#](#_CPPv43maxRK5arrayib14StreamOrDevice "Link to this definition")
The maximum of the elements of an array along the given axis.
array min(const array &a, bool keepdims, StreamOrDevice s \= {})[#](#_CPPv43minRK5arrayb14StreamOrDevice "Link to this definition")
The minimum of all elements of the array.
inline array min(const array &a, StreamOrDevice s \= {})[#](#_CPPv43minRK5array14StreamOrDevice "Link to this definition")
array min(const array &a, const [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector &axes, bool keepdims \= false, StreamOrDevice s \= {})[#](#_CPPv43minRK5arrayRKNSt6vectorIiEEb14StreamOrDevice "Link to this definition")
The minimum of the elements of an array along the given axes.
array min(const array &a, int axis, bool keepdims \= false, StreamOrDevice s \= {})[#](#_CPPv43minRK5arrayib14StreamOrDevice "Link to this definition")
The minimum of the elements of an array along the given axis.
array argmin(const array &a, bool keepdims, StreamOrDevice s \= {})[#](#_CPPv46argminRK5arrayb14StreamOrDevice "Link to this definition")
Returns the index of the minimum value in the array.
inline array argmin(const array &a, StreamOrDevice s \= {})[#](#_CPPv46argminRK5array14StreamOrDevice "Link to this definition")
array argmin(const array &a, int axis, bool keepdims \= false, StreamOrDevice s \= {})[#](#_CPPv46argminRK5arrayib14StreamOrDevice "Link to this definition")
Returns the indices of the minimum values along a given axis.
array argmax(const array &a, bool keepdims, StreamOrDevice s \= {})[#](#_CPPv46argmaxRK5arrayb14StreamOrDevice "Link to this definition")
Returns the index of the maximum value in the array.
inline array argmax(const array &a, StreamOrDevice s \= {})[#](#_CPPv46argmaxRK5array14StreamOrDevice "Link to this definition")
array argmax(const array &a, int axis, bool keepdims \= false, StreamOrDevice s \= {})[#](#_CPPv46argmaxRK5arrayib14StreamOrDevice "Link to this definition")
Returns the indices of the maximum values along a given axis.
array sort(const array &a, StreamOrDevice s \= {})[#](#_CPPv44sortRK5array14StreamOrDevice "Link to this definition")
Returns a sorted copy of the flattened array.
array sort(const array &a, int axis, StreamOrDevice s \= {})[#](#_CPPv44sortRK5arrayi14StreamOrDevice "Link to this definition")
Returns a sorted copy of the array along a given axis.
array argsort(const array &a, StreamOrDevice s \= {})[#](#_CPPv47argsortRK5array14StreamOrDevice "Link to this definition")
Returns indices that sort the flattened array.
array argsort(const array &a, int axis, StreamOrDevice s \= {})[#](#_CPPv47argsortRK5arrayi14StreamOrDevice "Link to this definition")
Returns indices that sort the array along a given axis.
array partition(const array &a, int kth, StreamOrDevice s \= {})[#](#_CPPv49partitionRK5arrayi14StreamOrDevice "Link to this definition")
Returns a partitioned copy of the flattened array such that the smaller kth elements are first.
array partition(const array &a, int kth, int axis, StreamOrDevice s \= {})[#](#_CPPv49partitionRK5arrayii14StreamOrDevice "Link to this definition")
Returns a partitioned copy of the array along a given axis such that the smaller kth elements are first.
array argpartition(const array &a, int kth, StreamOrDevice s \= {})[#](#_CPPv412argpartitionRK5arrayi14StreamOrDevice "Link to this definition")
Returns indices that partition the flattened array such that the smaller kth elements are first.
array argpartition(const array &a, int kth, int axis, StreamOrDevice s \= {})[#](#_CPPv412argpartitionRK5arrayii14StreamOrDevice "Link to this definition")
Returns indices that partition the array along a given axis such that the smaller kth elements are first.
array topk(const array &a, int k, StreamOrDevice s \= {})[#](#_CPPv44topkRK5arrayi14StreamOrDevice "Link to this definition")
Returns topk elements of the flattened array.
array topk(const array &a, int k, int axis, StreamOrDevice s \= {})[#](#_CPPv44topkRK5arrayii14StreamOrDevice "Link to this definition")
Returns topk elements of the array along a given axis.
array logsumexp(const array &a, bool keepdims, StreamOrDevice s \= {})[#](#_CPPv49logsumexpRK5arrayb14StreamOrDevice "Link to this definition")
The logsumexp of all elements of the array.
inline array logsumexp(const array &a, StreamOrDevice s \= {})[#](#_CPPv49logsumexpRK5array14StreamOrDevice "Link to this definition")
array logsumexp(const array &a, const [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector &axes, bool keepdims \= false, StreamOrDevice s \= {})[#](#_CPPv49logsumexpRK5arrayRKNSt6vectorIiEEb14StreamOrDevice "Link to this definition")
The logsumexp of the elements of an array along the given axes.
array logsumexp(const array &a, int axis, bool keepdims \= false, StreamOrDevice s \= {})[#](#_CPPv49logsumexpRK5arrayib14StreamOrDevice "Link to this definition")
The logsumexp of the elements of an array along the given axis.
array abs(const array &a, StreamOrDevice s \= {})[#](#_CPPv43absRK5array14StreamOrDevice "Link to this definition")
Absolute value of elements in an array.
array negative(const array &a, StreamOrDevice s \= {})[#](#_CPPv48negativeRK5array14StreamOrDevice "Link to this definition")
Negate an array.
array operator\-(const array &a)[#](#_CPPv4miRK5array "Link to this definition")
array sign(const array &a, StreamOrDevice s \= {})[#](#_CPPv44signRK5array14StreamOrDevice "Link to this definition")
The sign of the elements in an array.
array logical\_not(const array &a, StreamOrDevice s \= {})[#](#_CPPv411logical_notRK5array14StreamOrDevice "Link to this definition")
Logical not of an array.
array logical\_and(const array &a, const array &b, StreamOrDevice s \= {})[#](#_CPPv411logical_andRK5arrayRK5array14StreamOrDevice "Link to this definition")
Logical and of two arrays.
array operator&&(const array &a, const array &b)[#](#_CPPv4aaRK5arrayRK5array "Link to this definition")
array logical\_or(const array &a, const array &b, StreamOrDevice s \= {})[#](#_CPPv410logical_orRK5arrayRK5array14StreamOrDevice "Link to this definition")
Logical or of two arrays.
array operator||(const array &a, const array &b)[#](#_CPPv4ooRK5arrayRK5array "Link to this definition")
array reciprocal(const array &a, StreamOrDevice s \= {})[#](#_CPPv410reciprocalRK5array14StreamOrDevice "Link to this definition")
The reciprocal (1/x) of the elements in an array.
array add(const array &a, const array &b, StreamOrDevice s \= {})[#](#_CPPv43addRK5arrayRK5array14StreamOrDevice "Link to this definition")
Add two arrays.
array operator+(const array &a, const array &b)[#](#_CPPv4plRK5arrayRK5array "Link to this definition")
template
array operator+([T](#_CPPv4I0Epl5array1TRK5array "operator+::T")
a, const array &b)[#](#_CPPv4I0Epl5array1TRK5array "Link to this definition")
template
array operator+(const array &a, [T](#_CPPv4I0Epl5arrayRK5array1T "operator+::T")
b)[#](#_CPPv4I0Epl5arrayRK5array1T "Link to this definition")
array subtract(const array &a, const array &b, StreamOrDevice s \= {})[#](#_CPPv48subtractRK5arrayRK5array14StreamOrDevice "Link to this definition")
Subtract two arrays.
array operator\-(const array &a, const array &b)[#](#_CPPv4miRK5arrayRK5array "Link to this definition")
template
array operator\-([T](#_CPPv4I0Emi5array1TRK5array "operator-::T")
a, const array &b)[#](#_CPPv4I0Emi5array1TRK5array "Link to this definition")
template
array operator\-(const array &a, [T](#_CPPv4I0Emi5arrayRK5array1T "operator-::T")
b)[#](#_CPPv4I0Emi5arrayRK5array1T "Link to this definition")
array multiply(const array &a, const array &b, StreamOrDevice s \= {})[#](#_CPPv48multiplyRK5arrayRK5array14StreamOrDevice "Link to this definition")
Multiply two arrays.
array operator\*(const array &a, const array &b)[#](#_CPPv4mlRK5arrayRK5array "Link to this definition")
template
array operator\*([T](#_CPPv4I0Eml5array1TRK5array "operator*::T")
a, const array &b)[#](#_CPPv4I0Eml5array1TRK5array "Link to this definition")
template
array operator\*(const array &a, [T](#_CPPv4I0Eml5arrayRK5array1T "operator*::T")
b)[#](#_CPPv4I0Eml5arrayRK5array1T "Link to this definition")
array divide(const array &a, const array &b, StreamOrDevice s \= {})[#](#_CPPv46divideRK5arrayRK5array14StreamOrDevice "Link to this definition")
Divide two arrays.
array operator/(const array &a, const array &b)[#](#_CPPv4dvRK5arrayRK5array "Link to this definition")
array operator/(double a, const array &b)[#](#_CPPv4dvdRK5array "Link to this definition")
array operator/(const array &a, double b)[#](#_CPPv4dvRK5arrayd "Link to this definition")
[std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector divmod(const array &a, const array &b, StreamOrDevice s \= {})[#](#_CPPv46divmodRK5arrayRK5array14StreamOrDevice "Link to this definition")
Compute the element-wise quotient and remainder.
array floor\_divide(const array &a, const array &b, StreamOrDevice s \= {})[#](#_CPPv412floor_divideRK5arrayRK5array14StreamOrDevice "Link to this definition")
Compute integer division.
Equivalent to doing floor(a / x).
array remainder(const array &a, const array &b, StreamOrDevice s \= {})[#](#_CPPv49remainderRK5arrayRK5array14StreamOrDevice "Link to this definition")
Compute the element-wise remainder of division.
array operator%(const array &a, const array &b)[#](#_CPPv4rmRK5arrayRK5array "Link to this definition")
template
array operator%([T](#_CPPv4I0Erm5array1TRK5array "operator%::T")
a, const array &b)[#](#_CPPv4I0Erm5array1TRK5array "Link to this definition")
template
array operator%(const array &a, [T](#_CPPv4I0Erm5arrayRK5array1T "operator%::T")
b)[#](#_CPPv4I0Erm5arrayRK5array1T "Link to this definition")
array maximum(const array &a, const array &b, StreamOrDevice s \= {})[#](#_CPPv47maximumRK5arrayRK5array14StreamOrDevice "Link to this definition")
Element-wise maximum between two arrays.
array minimum(const array &a, const array &b, StreamOrDevice s \= {})[#](#_CPPv47minimumRK5arrayRK5array14StreamOrDevice "Link to this definition")
Element-wise minimum between two arrays.
array floor(const array &a, StreamOrDevice s \= {})[#](#_CPPv45floorRK5array14StreamOrDevice "Link to this definition")
Floor the element of an array.
array ceil(const array &a, StreamOrDevice s \= {})[#](#_CPPv44ceilRK5array14StreamOrDevice "Link to this definition")
Ceil the element of an array.
array square(const array &a, StreamOrDevice s \= {})[#](#_CPPv46squareRK5array14StreamOrDevice "Link to this definition")
Square the elements of an array.
array exp(const array &a, StreamOrDevice s \= {})[#](#_CPPv43expRK5array14StreamOrDevice "Link to this definition")
Exponential of the elements of an array.
array sin(const array &a, StreamOrDevice s \= {})[#](#_CPPv43sinRK5array14StreamOrDevice "Link to this definition")
Sine of the elements of an array.
array cos(const array &a, StreamOrDevice s \= {})[#](#_CPPv43cosRK5array14StreamOrDevice "Link to this definition")
Cosine of the elements of an array.
array tan(const array &a, StreamOrDevice s \= {})[#](#_CPPv43tanRK5array14StreamOrDevice "Link to this definition")
Tangent of the elements of an array.
array arcsin(const array &a, StreamOrDevice s \= {})[#](#_CPPv46arcsinRK5array14StreamOrDevice "Link to this definition")
Arc Sine of the elements of an array.
array arccos(const array &a, StreamOrDevice s \= {})[#](#_CPPv46arccosRK5array14StreamOrDevice "Link to this definition")
Arc Cosine of the elements of an array.
array arctan(const array &a, StreamOrDevice s \= {})[#](#_CPPv46arctanRK5array14StreamOrDevice "Link to this definition")
Arc Tangent of the elements of an array.
array arctan2(const array &a, const array &b, StreamOrDevice s \= {})[#](#_CPPv47arctan2RK5arrayRK5array14StreamOrDevice "Link to this definition")
Inverse tangent of the ratio of two arrays.
array sinh(const array &a, StreamOrDevice s \= {})[#](#_CPPv44sinhRK5array14StreamOrDevice "Link to this definition")
Hyperbolic Sine of the elements of an array.
array cosh(const array &a, StreamOrDevice s \= {})[#](#_CPPv44coshRK5array14StreamOrDevice "Link to this definition")
Hyperbolic Cosine of the elements of an array.
array tanh(const array &a, StreamOrDevice s \= {})[#](#_CPPv44tanhRK5array14StreamOrDevice "Link to this definition")
Hyperbolic Tangent of the elements of an array.
array arcsinh(const array &a, StreamOrDevice s \= {})[#](#_CPPv47arcsinhRK5array14StreamOrDevice "Link to this definition")
Inverse Hyperbolic Sine of the elements of an array.
array arccosh(const array &a, StreamOrDevice s \= {})[#](#_CPPv47arccoshRK5array14StreamOrDevice "Link to this definition")
Inverse Hyperbolic Cosine of the elements of an array.
array arctanh(const array &a, StreamOrDevice s \= {})[#](#_CPPv47arctanhRK5array14StreamOrDevice "Link to this definition")
Inverse Hyperbolic Tangent of the elements of an array.
array degrees(const array &a, StreamOrDevice s \= {})[#](#_CPPv47degreesRK5array14StreamOrDevice "Link to this definition")
Convert the elements of an array from Radians to Degrees.
array radians(const array &a, StreamOrDevice s \= {})[#](#_CPPv47radiansRK5array14StreamOrDevice "Link to this definition")
Convert the elements of an array from Degrees to Radians.
array log(const array &a, StreamOrDevice s \= {})[#](#_CPPv43logRK5array14StreamOrDevice "Link to this definition")
Natural logarithm of the elements of an array.
array log2(const array &a, StreamOrDevice s \= {})[#](#_CPPv44log2RK5array14StreamOrDevice "Link to this definition")
Log base 2 of the elements of an array.
array log10(const array &a, StreamOrDevice s \= {})[#](#_CPPv45log10RK5array14StreamOrDevice "Link to this definition")
Log base 10 of the elements of an array.
array log1p(const array &a, StreamOrDevice s \= {})[#](#_CPPv45log1pRK5array14StreamOrDevice "Link to this definition")
Natural logarithm of one plus elements in the array: `log(1 + a)`.
array logaddexp(const array &a, const array &b, StreamOrDevice s \= {})[#](#_CPPv49logaddexpRK5arrayRK5array14StreamOrDevice "Link to this definition")
Log-add-exp of one elements in the array: `log(exp(a) + exp(b))`.
array sigmoid(const array &a, StreamOrDevice s \= {})[#](#_CPPv47sigmoidRK5array14StreamOrDevice "Link to this definition")
Element-wise logistic sigmoid of the array: `1 / (1 + exp(-x)`.
array erf(const array &a, StreamOrDevice s \= {})[#](#_CPPv43erfRK5array14StreamOrDevice "Link to this definition")
Computes the error function of the elements of an array.
array erfinv(const array &a, StreamOrDevice s \= {})[#](#_CPPv46erfinvRK5array14StreamOrDevice "Link to this definition")
Computes the inverse error function of the elements of an array.
array expm1(const array &a, StreamOrDevice s \= {})[#](#_CPPv45expm1RK5array14StreamOrDevice "Link to this definition")
Computes the expm1 function of the elements of an array.
array stop\_gradient(const array &a, StreamOrDevice s \= {})[#](#_CPPv413stop_gradientRK5array14StreamOrDevice "Link to this definition")
Stop the flow of gradients.
array round(const array &a, int decimals, StreamOrDevice s \= {})[#](#_CPPv45roundRK5arrayi14StreamOrDevice "Link to this definition")
Round a floating point number.
inline array round(const array &a, StreamOrDevice s \= {})[#](#_CPPv45roundRK5array14StreamOrDevice "Link to this definition")
array matmul(const array &a, const array &b, StreamOrDevice s \= {})[#](#_CPPv46matmulRK5arrayRK5array14StreamOrDevice "Link to this definition")
Matrix-matrix multiplication.
array gather(const array &a, const [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector &indices, const [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector &axes, const Shape &slice\_sizes, StreamOrDevice s \= {})[#](#_CPPv46gatherRK5arrayRKNSt6vectorI5arrayEERKNSt6vectorIiEERK5Shape14StreamOrDevice "Link to this definition")
Gather array entries given indices and slices.
inline array gather(const array &a, const array &indices, int axis, const Shape &slice\_sizes, StreamOrDevice s \= {})[#](#_CPPv46gatherRK5arrayRK5arrayiRK5Shape14StreamOrDevice "Link to this definition")
array kron(const array &a, const array &b, StreamOrDevice s \= {})[#](#_CPPv44kronRK5arrayRK5array14StreamOrDevice "Link to this definition")
Compute the Kronecker product of two arrays.
array take(const array &a, const array &indices, int axis, StreamOrDevice s \= {})[#](#_CPPv44takeRK5arrayRK5arrayi14StreamOrDevice "Link to this definition")
Take array slices at the given indices of the specified axis.
array take(const array &a, int index, int axis, StreamOrDevice s \= {})[#](#_CPPv44takeRK5arrayii14StreamOrDevice "Link to this definition")
array take(const array &a, const array &indices, StreamOrDevice s \= {})[#](#_CPPv44takeRK5arrayRK5array14StreamOrDevice "Link to this definition")
Take array entries at the given indices treating the array as flattened.
array take(const array &a, int index, StreamOrDevice s \= {})[#](#_CPPv44takeRK5arrayi14StreamOrDevice "Link to this definition")
array take\_along\_axis(const array &a, const array &indices, int axis, StreamOrDevice s \= {})[#](#_CPPv415take_along_axisRK5arrayRK5arrayi14StreamOrDevice "Link to this definition")
Take array entries given indices along the axis.
array put\_along\_axis(const array &a, const array &indices, const array &values, int axis, StreamOrDevice s \= {})[#](#_CPPv414put_along_axisRK5arrayRK5arrayRK5arrayi14StreamOrDevice "Link to this definition")
Put the values into the array at the given indices along the axis.
array scatter\_add\_axis(const array &a, const array &indices, const array &values, int axis, StreamOrDevice s \= {})[#](#_CPPv416scatter_add_axisRK5arrayRK5arrayRK5arrayi14StreamOrDevice "Link to this definition")
Add the values into the array at the given indices along the axis.
array scatter(const array &a, const [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector &indices, const array &updates, const [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector &axes, StreamOrDevice s \= {})[#](#_CPPv47scatterRK5arrayRKNSt6vectorI5arrayEERK5arrayRKNSt6vectorIiEE14StreamOrDevice "Link to this definition")
Scatter updates to the given indices.
The parameters `indices` and `axes` determine the locations of `a` that are updated with the values in `updates`. Assuming 1-d `indices` for simplicity, `indices[i]` are the indices on axis `axes[i]` to which the values in `updates` will be applied. Note each array in `indices` is assigned to a corresponding axis and hence `indices.size() == axes.size()`. If an index/axis pair is not provided then indices along that axis are assumed to be zero.
Note the rank of `updates` must be equal to the sum of the rank of the broadcasted `indices` and the rank of `a`. In other words, assuming the arrays in `indices` have the same shape, `updates.ndim() == indices[0].ndim() + a.ndim()`. The leading dimensions of `updates` correspond to the indices, and the remaining `a.ndim()` dimensions are the values that will be applied to the given location in `a`.
For example:
auto in \= zeros({4, 4}, float32);
auto indices \= array({2});
auto updates \= reshape(arange(1, 3, float32), {1, 1, 2});
std::vector axes{0};
auto out \= scatter(in, {indices}, updates, axes);
will produce:
array(\[\[0, 0, 0, 0\],\
\[0, 0, 0, 0\],\
\[1, 2, 0, 0\],\
\[0, 0, 0, 0\]\], dtype\=float32)
This scatters the two-element row vector `[1, 2]` starting at the `(2, 0)` position of `a`.
Adding another element to `indices` will scatter into another location of `a`. We also have to add an another update for the new index:
auto in \= zeros({4, 4}, float32);
auto indices \= array({2, 0});
auto updates \= reshape(arange(1, 5, float32), {2, 1, 2});
std::vector axes{0};
auto out \= scatter(in, {indices}, updates, axes):
will produce:
array(\[\[3, 4, 0, 0\],\
\[0, 0, 0, 0\],\
\[1, 2, 0, 0\],\
\[0, 0, 0, 0\]\], dtype\=float32)
To control the scatter location on an additional axis, add another index array to `indices` and another axis to `axes`:
auto in \= zeros({4, 4}, float32);
auto indices \= std::vector{array({2, 0}), array({1, 2})};
auto updates \= reshape(arange(1, 5, float32), {2, 1, 2});
std::vector axes{0, 1};
auto out \= scatter(in, indices, updates, axes);
will produce:
array(\[\[0, 0, 3, 4\],\
\[0, 0, 0, 0\],\
\[0, 1, 2, 0\],\
\[0, 0, 0, 0\]\], dtype\=float32)
Items in indices are broadcasted together. This means:
auto indices \= std::vector{array({2, 0}), array({1})};
is equivalent to:
auto indices \= std::vector{array({2, 0}), array({1, 1})};
Note, `scatter` does not perform bounds checking on the indices and updates. Out-of-bounds accesses on `a` are undefined and typically result in unintended or invalid memory writes.
inline array scatter(const array &a, const array &indices, const array &updates, int axis, StreamOrDevice s \= {})[#](#_CPPv47scatterRK5arrayRK5arrayRK5arrayi14StreamOrDevice "Link to this definition")
array scatter\_add(const array &a, const [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector &indices, const array &updates, const [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector &axes, StreamOrDevice s \= {})[#](#_CPPv411scatter_addRK5arrayRKNSt6vectorI5arrayEERK5arrayRKNSt6vectorIiEE14StreamOrDevice "Link to this definition")
Scatter and add updates to given indices.
inline array scatter\_add(const array &a, const array &indices, const array &updates, int axis, StreamOrDevice s \= {})[#](#_CPPv411scatter_addRK5arrayRK5arrayRK5arrayi14StreamOrDevice "Link to this definition")
array scatter\_prod(const array &a, const [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector &indices, const array &updates, const [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector &axes, StreamOrDevice s \= {})[#](#_CPPv412scatter_prodRK5arrayRKNSt6vectorI5arrayEERK5arrayRKNSt6vectorIiEE14StreamOrDevice "Link to this definition")
Scatter and prod updates to given indices.
inline array scatter\_prod(const array &a, const array &indices, const array &updates, int axis, StreamOrDevice s \= {})[#](#_CPPv412scatter_prodRK5arrayRK5arrayRK5arrayi14StreamOrDevice "Link to this definition")
array scatter\_max(const array &a, const [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector &indices, const array &updates, const [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector &axes, StreamOrDevice s \= {})[#](#_CPPv411scatter_maxRK5arrayRKNSt6vectorI5arrayEERK5arrayRKNSt6vectorIiEE14StreamOrDevice "Link to this definition")
Scatter and max updates to given linear indices.
inline array scatter\_max(const array &a, const array &indices, const array &updates, int axis, StreamOrDevice s \= {})[#](#_CPPv411scatter_maxRK5arrayRK5arrayRK5arrayi14StreamOrDevice "Link to this definition")
array scatter\_min(const array &a, const [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector &indices, const array &updates, const [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector &axes, StreamOrDevice s \= {})[#](#_CPPv411scatter_minRK5arrayRKNSt6vectorI5arrayEERK5arrayRKNSt6vectorIiEE14StreamOrDevice "Link to this definition")
Scatter and min updates to given linear indices.
inline array scatter\_min(const array &a, const array &indices, const array &updates, int axis, StreamOrDevice s \= {})[#](#_CPPv411scatter_minRK5arrayRK5arrayRK5arrayi14StreamOrDevice "Link to this definition")
array sqrt(const array &a, StreamOrDevice s \= {})[#](#_CPPv44sqrtRK5array14StreamOrDevice "Link to this definition")
Square root the elements of an array.
array rsqrt(const array &a, StreamOrDevice s \= {})[#](#_CPPv45rsqrtRK5array14StreamOrDevice "Link to this definition")
Square root and reciprocal the elements of an array.
array softmax(const array &a, const [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector &axes, bool precise \= false, StreamOrDevice s \= {})[#](#_CPPv47softmaxRK5arrayRKNSt6vectorIiEEb14StreamOrDevice "Link to this definition")
Softmax of an array.
array softmax(const array &a, bool precise \= false, StreamOrDevice s \= {})[#](#_CPPv47softmaxRK5arrayb14StreamOrDevice "Link to this definition")
Softmax of an array.
inline array softmax(const array &a, int axis, bool precise \= false, StreamOrDevice s \= {})[#](#_CPPv47softmaxRK5arrayib14StreamOrDevice "Link to this definition")
Softmax of an array.
array power(const array &a, const array &b, StreamOrDevice s \= {})[#](#_CPPv45powerRK5arrayRK5array14StreamOrDevice "Link to this definition")
Raise elements of a to the power of b element-wise.
array cumsum(const array &a, int axis, bool reverse \= false, bool inclusive \= true, StreamOrDevice s \= {})[#](#_CPPv46cumsumRK5arrayibb14StreamOrDevice "Link to this definition")
Cumulative sum of an array.
array cumprod(const array &a, int axis, bool reverse \= false, bool inclusive \= true, StreamOrDevice s \= {})[#](#_CPPv47cumprodRK5arrayibb14StreamOrDevice "Link to this definition")
Cumulative product of an array.
array cummax(const array &a, int axis, bool reverse \= false, bool inclusive \= true, StreamOrDevice s \= {})[#](#_CPPv46cummaxRK5arrayibb14StreamOrDevice "Link to this definition")
Cumulative max of an array.
array cummin(const array &a, int axis, bool reverse \= false, bool inclusive \= true, StreamOrDevice s \= {})[#](#_CPPv46cumminRK5arrayibb14StreamOrDevice "Link to this definition")
Cumulative min of an array.
array conv\_general(array input, array weight, [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector stride \= {}, [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector padding\_lo \= {}, [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector padding\_hi \= {}, [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector kernel\_dilation \= {}, [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector input\_dilation \= {}, int groups \= 1, bool flip \= false, StreamOrDevice s \= {})[#](#_CPPv412conv_general5array5arrayNSt6vectorIiEENSt6vectorIiEENSt6vectorIiEENSt6vectorIiEENSt6vectorIiEEib14StreamOrDevice "Link to this definition")
General convolution with a filter.
inline array conv\_general(const array &input, const array &weight, [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector stride \= {}, [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector padding \= {}, [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector kernel\_dilation \= {}, [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector input\_dilation \= {}, int groups \= 1, bool flip \= false, StreamOrDevice s \= {})[#](#_CPPv412conv_generalRK5arrayRK5arrayNSt6vectorIiEENSt6vectorIiEENSt6vectorIiEENSt6vectorIiEEib14StreamOrDevice "Link to this definition")
General convolution with a filter.
array conv1d(const array &input, const array &weight, int stride \= 1, int padding \= 0, int dilation \= 1, int groups \= 1, StreamOrDevice s \= {})[#](#_CPPv46conv1dRK5arrayRK5arrayiiii14StreamOrDevice "Link to this definition")
1D convolution with a filter
array conv2d(const array &input, const array &weight, const [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::pair &stride \= {1, 1}, const [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::pair &padding \= {0, 0}, const [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::pair &dilation \= {1, 1}, int groups \= 1, StreamOrDevice s \= {})[#](#_CPPv46conv2dRK5arrayRK5arrayRKNSt4pairIiiEERKNSt4pairIiiEERKNSt4pairIiiEEi14StreamOrDevice "Link to this definition")
2D convolution with a filter
array conv3d(const array &input, const array &weight, const [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::tuple &stride \= {1, 1, 1}, const [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::tuple &padding \= {0, 0, 0}, const [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::tuple &dilation \= {1, 1, 1}, int groups \= 1, StreamOrDevice s \= {})[#](#_CPPv46conv3dRK5arrayRK5arrayRKNSt5tupleIiiiEERKNSt5tupleIiiiEERKNSt5tupleIiiiEEi14StreamOrDevice "Link to this definition")
3D convolution with a filter
array conv\_transpose1d(const array &input, const array &weight, int stride \= 1, int padding \= 0, int dilation \= 1, int groups \= 1, StreamOrDevice s \= {})[#](#_CPPv416conv_transpose1dRK5arrayRK5arrayiiii14StreamOrDevice "Link to this definition")
1D transposed convolution with a filter
array conv\_transpose2d(const array &input, const array &weight, const [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::pair &stride \= {1, 1}, const [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::pair &padding \= {0, 0}, const [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::pair &dilation \= {1, 1}, int groups \= 1, StreamOrDevice s \= {})[#](#_CPPv416conv_transpose2dRK5arrayRK5arrayRKNSt4pairIiiEERKNSt4pairIiiEERKNSt4pairIiiEEi14StreamOrDevice "Link to this definition")
2D transposed convolution with a filter
array conv\_transpose3d(const array &input, const array &weight, const [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::tuple &stride \= {1, 1, 1}, const [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::tuple &padding \= {0, 0, 0}, const [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::tuple &dilation \= {1, 1, 1}, int groups \= 1, StreamOrDevice s \= {})[#](#_CPPv416conv_transpose3dRK5arrayRK5arrayRKNSt5tupleIiiiEERKNSt5tupleIiiiEERKNSt5tupleIiiiEEi14StreamOrDevice "Link to this definition")
3D transposed convolution with a filter
array quantized\_matmul(array x, array w, array scales, array biases, bool transpose \= true, int group\_size \= 64, int bits \= 4, StreamOrDevice s \= {})[#](#_CPPv416quantized_matmul5array5array5array5arraybii14StreamOrDevice "Link to this definition")
Quantized matmul multiplies x with a quantized matrix w.
[std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::tuple quantize(const array &w, int group\_size \= 64, int bits \= 4, StreamOrDevice s \= {})[#](#_CPPv48quantizeRK5arrayii14StreamOrDevice "Link to this definition")
Quantize a matrix along its last axis.
array dequantize(const array &w, const array &scales, const array &biases, int group\_size \= 64, int bits \= 4, StreamOrDevice s \= {})[#](#_CPPv410dequantizeRK5arrayRK5arrayRK5arrayii14StreamOrDevice "Link to this definition")
Dequantize a matrix produced by [quantize()](#group__ops_1gab43cc28690da7cdd43b43065adbd31da)
array gather\_qmm(const array &x, const array &w, const array &scales, const array &biases, [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::optional lhs\_indices \= [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::nullopt, [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::optional rhs\_indices \= [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::nullopt, bool transpose \= true, int group\_size \= 64, int bits \= 4, StreamOrDevice s \= {})[#](#_CPPv410gather_qmmRK5arrayRK5arrayRK5arrayRK5arrayNSt8optionalI5arrayEENSt8optionalI5arrayEEbii14StreamOrDevice "Link to this definition")
Compute matrix products with matrix-level gather.
array tensordot(const array &a, const array &b, const int axis \= 2, StreamOrDevice s \= {})[#](#_CPPv49tensordotRK5arrayRK5arrayKi14StreamOrDevice "Link to this definition")
Returns a contraction of a and b over multiple dimensions.
array tensordot(const array &a, const array &b, const [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector &axes\_a, const [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector &axes\_b, StreamOrDevice s \= {})[#](#_CPPv49tensordotRK5arrayRK5arrayRKNSt6vectorIiEERKNSt6vectorIiEE14StreamOrDevice "Link to this definition")
array outer(const array &a, const array &b, StreamOrDevice s \= {})[#](#_CPPv45outerRK5arrayRK5array14StreamOrDevice "Link to this definition")
Compute the outer product of two vectors.
array inner(const array &a, const array &b, StreamOrDevice s \= {})[#](#_CPPv45innerRK5arrayRK5array14StreamOrDevice "Link to this definition")
Compute the inner product of two vectors.
array addmm(array c, array a, array b, const float &alpha \= 1.f, const float &beta \= 1.f, StreamOrDevice s \= {})[#](#_CPPv45addmm5array5array5arrayRKfRKf14StreamOrDevice "Link to this definition")
Compute D = beta \* C + alpha \* (A @ B)
array block\_masked\_mm(array a, array b, int block\_size, [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::optional mask\_out \= [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::nullopt, [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::optional mask\_lhs \= [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::nullopt, [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::optional mask\_rhs \= [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::nullopt, StreamOrDevice s \= {})[#](#_CPPv415block_masked_mm5array5arrayiNSt8optionalI5arrayEENSt8optionalI5arrayEENSt8optionalI5arrayEE14StreamOrDevice "Link to this definition")
Compute matrix product with block masking.
array gather\_mm(array a, array b, [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::optional lhs\_indices \= [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::nullopt, [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::optional rhs\_indices \= [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::nullopt, StreamOrDevice s \= {})[#](#_CPPv49gather_mm5array5arrayNSt8optionalI5arrayEENSt8optionalI5arrayEE14StreamOrDevice "Link to this definition")
Compute matrix product with matrix-level gather.
array diagonal(const array &a, int offset \= 0, int axis1 \= 0, int axis2 \= 1, StreamOrDevice s \= {})[#](#_CPPv48diagonalRK5arrayiii14StreamOrDevice "Link to this definition")
Extract a diagonal or construct a diagonal array.
array diag(const array &a, int k \= 0, StreamOrDevice s \= {})[#](#_CPPv44diagRK5arrayi14StreamOrDevice "Link to this definition")
Extract diagonal from a 2d array or create a diagonal matrix.
array trace(const array &a, int offset, int axis1, int axis2, Dtype dtype, StreamOrDevice s \= {})[#](#_CPPv45traceRK5arrayiii5Dtype14StreamOrDevice "Link to this definition")
Return the sum along a specified diagonal in the given array.
array trace(const array &a, int offset, int axis1, int axis2, StreamOrDevice s \= {})[#](#_CPPv45traceRK5arrayiii14StreamOrDevice "Link to this definition")
array trace(const array &a, StreamOrDevice s \= {})[#](#_CPPv45traceRK5array14StreamOrDevice "Link to this definition")
[std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector depends(const [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector &inputs, const [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector &dependencies)[#](#_CPPv47dependsRKNSt6vectorI5arrayEERKNSt6vectorI5arrayEE "Link to this definition")
Implements the identity function but allows injecting dependencies to other arrays.
This ensures that these other arrays will have been computed when the outputs of this function are computed.
array atleast\_1d(const array &a, StreamOrDevice s \= {})[#](#_CPPv410atleast_1dRK5array14StreamOrDevice "Link to this definition")
convert an array to an atleast ndim array
[std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector atleast\_1d(const [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector &a, StreamOrDevice s \= {})[#](#_CPPv410atleast_1dRKNSt6vectorI5arrayEE14StreamOrDevice "Link to this definition")
array atleast\_2d(const array &a, StreamOrDevice s \= {})[#](#_CPPv410atleast_2dRK5array14StreamOrDevice "Link to this definition")
[std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector atleast\_2d(const [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector &a, StreamOrDevice s \= {})[#](#_CPPv410atleast_2dRKNSt6vectorI5arrayEE14StreamOrDevice "Link to this definition")
array atleast\_3d(const array &a, StreamOrDevice s \= {})[#](#_CPPv410atleast_3dRK5array14StreamOrDevice "Link to this definition")
[std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector atleast\_3d(const [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector &a, StreamOrDevice s \= {})[#](#_CPPv410atleast_3dRKNSt6vectorI5arrayEE14StreamOrDevice "Link to this definition")
array number\_of\_elements(const array &a, [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector axes, bool inverted, Dtype dtype \= int32, StreamOrDevice s \= {})[#](#_CPPv418number_of_elementsRK5arrayNSt6vectorIiEEb5Dtype14StreamOrDevice "Link to this definition")
Extract the number of elements along some axes as a scalar array.
Used to allow shape dependent shapeless compilation (pun intended).
array conjugate(const array &a, StreamOrDevice s \= {})[#](#_CPPv49conjugateRK5array14StreamOrDevice "Link to this definition")
array bitwise\_and(const array &a, const array &b, StreamOrDevice s \= {})[#](#_CPPv411bitwise_andRK5arrayRK5array14StreamOrDevice "Link to this definition")
Bitwise and.
array operator&(const array &a, const array &b)[#](#_CPPv4anRK5arrayRK5array "Link to this definition")
array bitwise\_or(const array &a, const array &b, StreamOrDevice s \= {})[#](#_CPPv410bitwise_orRK5arrayRK5array14StreamOrDevice "Link to this definition")
Bitwise inclusive or.
array operator|(const array &a, const array &b)[#](#_CPPv4orRK5arrayRK5array "Link to this definition")
array bitwise\_xor(const array &a, const array &b, StreamOrDevice s \= {})[#](#_CPPv411bitwise_xorRK5arrayRK5array14StreamOrDevice "Link to this definition")
Bitwise exclusive or.
array operator^(const array &a, const array &b)[#](#_CPPv4eoRK5arrayRK5array "Link to this definition")
array left\_shift(const array &a, const array &b, StreamOrDevice s \= {})[#](#_CPPv410left_shiftRK5arrayRK5array14StreamOrDevice "Link to this definition")
Shift bits to the left.
array operator<<(const array &a, const array &b)[#](#_CPPv4lsRK5arrayRK5array "Link to this definition")
array right\_shift(const array &a, const array &b, StreamOrDevice s \= {})[#](#_CPPv411right_shiftRK5arrayRK5array14StreamOrDevice "Link to this definition")
Shift bits to the right.
array operator\>>(const array &a, const array &b)[#](#_CPPv4rsRK5arrayRK5array "Link to this definition")
array bitwise\_invert(const array &a, StreamOrDevice s \= {})[#](#_CPPv414bitwise_invertRK5array14StreamOrDevice "Link to this definition")
Invert the bits.
array operator~(const array &a)[#](#_CPPv4coRK5array "Link to this definition")
array view(const array &a, const Dtype &dtype, StreamOrDevice s \= {})[#](#_CPPv44viewRK5arrayRK5Dtype14StreamOrDevice "Link to this definition")
array roll(const array &a, int shift, StreamOrDevice s \= {})[#](#_CPPv44rollRK5arrayi14StreamOrDevice "Link to this definition")
Roll elements along an axis and introduce them on the other side.
array roll(const array &a, const Shape &shift, StreamOrDevice s \= {})[#](#_CPPv44rollRK5arrayRK5Shape14StreamOrDevice "Link to this definition")
array roll(const array &a, int shift, int axis, StreamOrDevice s \= {})[#](#_CPPv44rollRK5arrayii14StreamOrDevice "Link to this definition")
array roll(const array &a, int shift, const [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector &axes, StreamOrDevice s \= {})[#](#_CPPv44rollRK5arrayiRKNSt6vectorIiEE14StreamOrDevice "Link to this definition")
array roll(const array &a, const Shape &shift, int axis, StreamOrDevice s \= {})[#](#_CPPv44rollRK5arrayRK5Shapei14StreamOrDevice "Link to this definition")
array roll(const array &a, const Shape &shift, const [std](#_CPPv4StRK5arraybi14StreamOrDevice "std")
::vector &axes, StreamOrDevice s \= {})[#](#_CPPv44rollRK5arrayRK5ShapeRKNSt6vectorIiEE14StreamOrDevice "Link to this definition")
array real(const array &a, StreamOrDevice s \= {})[#](#_CPPv44realRK5array14StreamOrDevice "Link to this definition")
array imag(const array &a, StreamOrDevice s \= {})[#](#_CPPv44imagRK5array14StreamOrDevice "Link to this definition")
array contiguous(const array &a, bool allow\_col\_major \= false, StreamOrDevice s \= {})[#](#_CPPv410contiguousRK5arrayb14StreamOrDevice "Link to this definition")
Contents
---
# Custom Extensions in MLX — MLX 0.24.0 documentation
[Skip to main content](#main-content)
Back to top
Ctrl+K
[](../index.html)
[](https://github.com/ml-explore/mlx "Source repository")
* [.rst](../_sources/dev/extensions.rst "Download source file")
* .pdf
Custom Extensions in MLX
========================
Contents
--------
Custom Extensions in MLX[#](#custom-extensions-in-mlx "Link to this heading")
==============================================================================
You can extend MLX with custom operations on the CPU or GPU. This guide explains how to do that with a simple example.
Introducing the Example[#](#introducing-the-example "Link to this heading")
----------------------------------------------------------------------------
Let’s say you would like an operation that takes in two arrays, `x` and `y`, scales them both by coefficients `alpha` and `beta` respectively, and then adds them together to get the result `z = alpha * x + beta * y`. You can do that in MLX directly:
import mlx.core as mx
def simple\_axpby(x: mx.array, y: mx.array, alpha: float, beta: float) \-> mx.array:
return alpha \* x + beta \* y
This function performs that operation while leaving the implementation and function transformations to MLX.
However, you may want to customize the underlying implementation, perhaps to make it faster. In this tutorial we will go through adding custom extensions. It will cover:
* The structure of the MLX library.
* Implementing a CPU operation.
* Implementing a GPU operation using metal.
* Adding the `vjp` and `jvp` function transformation.
* Building a custom extension and binding it to python.
Operations and Primitives[#](#operations-and-primitives "Link to this heading")
--------------------------------------------------------------------------------
Operations in MLX build the computation graph. Primitives provide the rules for evaluating and transforming the graph. Let’s start by discussing operations in more detail.
### Operations[#](#operations "Link to this heading")
Operations are the front-end functions that operate on arrays. They are defined in the C++ API ([Operations](../cpp/ops.html#cpp-ops)
), and the Python API ([Operations](../python/ops.html#ops)
) binds them.
We would like an operation `axpby()` that takes in two arrays, `x` and `y`, and two scalars, `alpha` and `beta`. This is how to define it in C++:
/\*\*
\* Scale and sum two vectors element-wise
\* z = alpha \* x + beta \* y
\*
\* Use NumPy-style broadcasting between x and y
\* Inputs are upcasted to floats if needed
\*\*/
array axpby(
const array& x, // Input array x
const array& y, // Input array y
const float alpha, // Scaling factor for x
const float beta, // Scaling factor for y
StreamOrDevice s \= {} // Stream on which to schedule the operation
);
The simplest way to implement this is with existing operations:
array axpby(
const array& x, // Input array x
const array& y, // Input array y
const float alpha, // Scaling factor for x
const float beta, // Scaling factor for y
StreamOrDevice s /\* = {} \*/ // Stream on which to schedule the operation
) {
// Scale x and y on the provided stream
auto ax \= multiply(array(alpha), x, s);
auto by \= multiply(array(beta), y, s);
// Add and return
return add(ax, by, s);
}
The operations themselves do not contain the implementations that act on the data, nor do they contain the rules of transformations. Rather, they are an easy to use interface that use `Primitive` building blocks.
### Primitives[#](#primitives "Link to this heading")
A `Primitive` is part of the computation graph of an `array`. It defines how to create outputs arrays given a input arrays. Further, a `Primitive` has methods to run on the CPU or GPU and for function transformations such as `vjp` and `jvp`. Lets go back to our example to be more concrete:
class Axpby : public Primitive {
public:
explicit Axpby(Stream stream, float alpha, float beta)
: Primitive(stream), alpha\_(alpha), beta\_(beta){};
/\*\*
\* A primitive must know how to evaluate itself on the CPU/GPU
\* for the given inputs and populate the output array.
\*
\* To avoid unnecessary allocations, the evaluation function
\* is responsible for allocating space for the array.
\*/
void eval\_cpu(
const std::vector& inputs,
std::vector& outputs) override;
void eval\_gpu(
const std::vector& inputs,
std::vector& outputs) override;
/\*\* The Jacobian-vector product. \*/
std::vector jvp(
const std::vector& primals,
const std::vector& tangents,
const std::vector& argnums) override;
/\*\* The vector-Jacobian product. \*/
std::vector vjp(
const std::vector& primals,
const array& cotan,
const std::vector& argnums,
const std::vector& outputs) override;
/\*\*
\* The primitive must know how to vectorize itself across
\* the given axes. The output is a pair containing the array
\* representing the vectorized computation and the axis which
\* corresponds to the output vectorized dimension.
\*/
virtual std::pair, std::vector> vmap(
const std::vector& inputs,
const std::vector& axes) override;
/\*\* Print the primitive. \*/
void print(std::ostream& os) override {
os << "Axpby";
}
/\*\* Equivalence check \*\*/
bool is\_equivalent(const Primitive& other) const override;
private:
float alpha\_;
float beta\_;
};
The `Axpby` class derives from the base `Primitive` class. The `Axpby` treats `alpha` and `beta` as parameters. It then provides implementations of how the output array is produced given the inputs through `Axpby::eval_cpu()` and `Axpby::eval_gpu()`. It also provides rules of transformations in `Axpby::jvp()`, `Axpby::vjp()`, and `Axpby::vmap()`.
### Using the Primitive[#](#using-the-primitive "Link to this heading")
Operations can use this `Primitive` to add a new `array` to the computation graph. An `array` can be constructed by providing its data type, shape, the `Primitive` that computes it, and the `array` inputs that are passed to the primitive.
Let’s reimplement our operation now in terms of our `Axpby` primitive.
array axpby(
const array& x, // Input array x
const array& y, // Input array y
const float alpha, // Scaling factor for x
const float beta, // Scaling factor for y
StreamOrDevice s /\* = {} \*/ // Stream on which to schedule the operation
) {
// Promote dtypes between x and y as needed
auto promoted\_dtype \= promote\_types(x.dtype(), y.dtype());
// Upcast to float32 for non-floating point inputs x and y
auto out\_dtype \= issubdtype(promoted\_dtype, float32)
? promoted\_dtype
: promote\_types(promoted\_dtype, float32);
// Cast x and y up to the determined dtype (on the same stream s)
auto x\_casted \= astype(x, out\_dtype, s);
auto y\_casted \= astype(y, out\_dtype, s);
// Broadcast the shapes of x and y (on the same stream s)
auto broadcasted\_inputs \= broadcast\_arrays({x\_casted, y\_casted}, s);
auto out\_shape \= broadcasted\_inputs\[0\].shape();
// Construct the array as the output of the Axpby primitive
// with the broadcasted and upcasted arrays as inputs
return array(
/\* const std::vector& shape = \*/ out\_shape,
/\* Dtype dtype = \*/ out\_dtype,
/\* std::unique\_ptr primitive = \*/
std::make\_shared(to\_stream(s), alpha, beta),
/\* const std::vector& inputs = \*/ broadcasted\_inputs);
}
This operation now handles the following:
1. Upcast inputs and resolve the output data type.
2. Broadcast the inputs and resolve the output shape.
3. Construct the primitive `Axpby` using the given stream, `alpha`, and `beta`.
4. Construct the output `array` using the primitive and the inputs.
Implementing the Primitive[#](#implementing-the-primitive "Link to this heading")
----------------------------------------------------------------------------------
No computation happens when we call the operation alone. The operation only builds the computation graph. When we evaluate the output array, MLX schedules the execution of the computation graph, and calls `Axpby::eval_cpu()` or `Axpby::eval_gpu()` depending on the stream/device specified by the user.
Warning
When `Primitive::eval_cpu()` or `Primitive::eval_gpu()` are called, no memory has been allocated for the output array. It falls on the implementation of these functions to allocate memory as needed.
### Implementing the CPU Back-end[#](#implementing-the-cpu-back-end "Link to this heading")
Let’s start by implementing `Axpby::eval_cpu()`.
The method will go over each element of the output array, find the corresponding input elements of `x` and `y` and perform the operation point-wise. This is captured in the templated function `axpby_impl()`.
template
void axpby\_impl(
const mx::array& x,
const mx::array& y,
mx::array& out,
float alpha\_,
float beta\_,
mx::Stream stream) {
// Allocate the output with \`malloc\_or\_wait\` which synchronously allocates
// memory, potentially waiting if the system is under memory pressure
out.set\_data(mx::allocator::malloc\_or\_wait(out.nbytes()));
// Get the CPU command encoder and register input and output arrays
auto& encoder \= mx::cpu::get\_command\_encoder(stream);
encoder.set\_input\_array(x);
encoder.set\_input\_array(y);
encoder.set\_output\_array(out);
// Launch the CPU kernel
encoder.dispatch(\[x\_ptr \= x.data(),\
y\_ptr \= y.data(),\
out\_ptr \= out.data(),\
size \= out.size(),\
shape \= out.shape(),\
x\_strides \= x.strides(),\
y\_strides \= y.strides(),\
alpha\_,\
beta\_\]() {
// Cast alpha and beta to the relevant types
T alpha \= static\_cast(alpha\_);
T beta \= static\_cast(beta\_);
// Do the element-wise operation for each output
for (size\_t out\_idx \= 0; out\_idx < size; out\_idx++) {
// Map linear indices to offsets in x and y
auto x\_offset \= mx::elem\_to\_loc(out\_idx, shape, x\_strides);
auto y\_offset \= mx::elem\_to\_loc(out\_idx, shape, y\_strides);
// We allocate the output to be contiguous and regularly strided
// (defaults to row major) and hence it doesn't need additional mapping
out\_ptr\[out\_idx\] \= alpha \* x\_ptr\[x\_offset\] + beta \* y\_ptr\[y\_offset\];
}
});
}
Our implementation should work for all incoming floating point arrays. Accordingly, we add dispatches for `float32`, `float16`, `bfloat16` and `complex64`. We throw an error if we encounter an unexpected type.
void Axpby::eval\_cpu(
const std::vector& inputs,
std::vector& outputs) {
auto& x \= inputs\[0\];
auto& y \= inputs\[1\];
auto& out \= outputs\[0\];
// Dispatch to the correct dtype
if (out.dtype() \== mx::float32) {
return axpby\_impl(x, y, out, alpha\_, beta\_, stream());
} else if (out.dtype() \== mx::float16) {
return axpby\_impl(x, y, out, alpha\_, beta\_, stream());
} else if (out.dtype() \== mx::bfloat16) {
return axpby\_impl(x, y, out, alpha\_, beta\_, stream());
} else if (out.dtype() \== mx::complex64) {
return axpby\_impl(x, y, out, alpha\_, beta\_, stream());
} else {
throw std::runtime\_error(
"Axpby is only supported for floating point types.");
}
}
Just this much is enough to run the operation `axpby()` on a CPU stream! If you do not plan on running the operation on the GPU or using transforms on computation graphs that contain `Axpby`, you can stop implementing the primitive here.
### Implementing the GPU Back-end[#](#implementing-the-gpu-back-end "Link to this heading")
Apple silicon devices address their GPUs using the [Metal](https://developer.apple.com/documentation/metal?language=objc)
shading language, and GPU kernels in MLX are written using Metal.
Note
Here are some helpful resources if you are new to Metal:
* A walkthrough of the metal compute pipeline: [Metal Example](https://developer.apple.com/documentation/metal/performing_calculations_on_a_gpu?language=objc)
* Documentation for metal shading language: [Metal Specification](https://developer.apple.com/metal/Metal-Shading-Language-Specification.pdf)
* Using metal from C++: [Metal-cpp](https://developer.apple.com/metal/cpp/)
Let’s keep the GPU kernel simple. We will launch exactly as many threads as there are elements in the output. Each thread will pick the element it needs from `x` and `y`, do the point-wise operation, and update its assigned element in the output.
template
\[\[kernel\]\] void axpby\_general(
device const T\* x \[\[buffer(0)\]\],
device const T\* y \[\[buffer(1)\]\],
device T\* out \[\[buffer(2)\]\],
constant const float& alpha \[\[buffer(3)\]\],
constant const float& beta \[\[buffer(4)\]\],
constant const int\* shape \[\[buffer(5)\]\],
constant const int64\_t\* x\_strides \[\[buffer(6)\]\],
constant const int64\_t\* y\_strides \[\[buffer(7)\]\],
constant const int& ndim \[\[buffer(8)\]\],
uint index \[\[thread\_position\_in\_grid\]\]) {
// Convert linear indices to offsets in array
auto x\_offset \= elem\_to\_loc(index, shape, x\_strides, ndim);
auto y\_offset \= elem\_to\_loc(index, shape, y\_strides, ndim);
// Do the operation and update the output
out\[index\] \=
static\_cast(alpha) \* x\[x\_offset\] + static\_cast(beta) \* y\[y\_offset\];
}
We then need to instantiate this template for all floating point types and give each instantiation a unique host name so we can identify it.
instantiate\_kernel("axpby\_general\_float32", axpby\_general, float)
instantiate\_kernel("axpby\_general\_float16", axpby\_general, float16\_t)
instantiate\_kernel("axpby\_general\_bfloat16", axpby\_general, bfloat16\_t)
instantiate\_kernel("axpby\_general\_complex64", axpby\_general, complex64\_t)
The logic to determine the kernel, set the inputs, resolve the grid dimensions, and dispatch to the GPU are contained in `Axpby::eval_gpu()` as shown below.
/\*\* Evaluate primitive on GPU \*/
void Axpby::eval\_gpu(
const std::vector& inputs,
std::vector& outputs) {
// Prepare inputs
assert(inputs.size() \== 2);
auto& x \= inputs\[0\];
auto& y \= inputs\[1\];
auto& out \= outputs\[0\];
// Each primitive carries the stream it should execute on
// and each stream carries its device identifiers
auto& s \= stream();
// We get the needed metal device using the stream
auto& d \= metal::device(s.device);
// Allocate output memory
out.set\_data(allocator::malloc\_or\_wait(out.nbytes()));
// Resolve name of kernel
std::ostringstream kname;
kname << "axpby\_" << "general\_" << type\_to\_name(out);
// Make sure the metal library is available
d.register\_library("mlx\_ext");
// Make a kernel from this metal library
auto kernel \= d.get\_kernel(kname.str(), "mlx\_ext");
// Prepare to encode kernel
auto& compute\_encoder \= d.get\_command\_encoder(s.index);
compute\_encoder.set\_compute\_pipeline\_state(kernel);
// Kernel parameters are registered with buffer indices corresponding to
// those in the kernel declaration at axpby.metal
int ndim \= out.ndim();
size\_t nelem \= out.size();
// Encode input arrays to kernel
compute\_encoder.set\_input\_array(x, 0);
compute\_encoder.set\_input\_array(y, 1);
// Encode output arrays to kernel
compute\_encoder.set\_output\_array(out, 2);
// Encode alpha and beta
compute\_encoder.set\_bytes(alpha\_, 3);
compute\_encoder.set\_bytes(beta\_, 4);
// Encode shape, strides and ndim
compute\_encoder.set\_vector\_bytes(x.shape(), 5);
compute\_encoder.set\_vector\_bytes(x.strides(), 6);
compute\_encoder.set\_bytes(y.strides(), 7);
compute\_encoder.set\_bytes(ndim, 8);
// We launch 1 thread for each input and make sure that the number of
// threads in any given threadgroup is not higher than the max allowed
size\_t tgp\_size \= std::min(nelem, kernel\->maxTotalThreadsPerThreadgroup());
// Fix the 3D size of each threadgroup (in terms of threads)
MTL::Size group\_dims \= MTL::Size(tgp\_size, 1, 1);
// Fix the 3D size of the launch grid (in terms of threads)
MTL::Size grid\_dims \= MTL::Size(nelem, 1, 1);
// Launch the grid with the given number of threads divided among
// the given threadgroups
compute\_encoder.dispatch\_threads(grid\_dims, group\_dims);
}
We can now call the `axpby()` operation on both the CPU and the GPU!
A few things to note about MLX and Metal before moving on. MLX keeps track of the active `command_buffer` and the `MTLCommandBuffer` to which it is associated. We rely on `d.get_command_encoder()` to give us the active metal compute command encoder instead of building a new one and calling `compute_encoder->end_encoding()` at the end. MLX adds kernels (compute pipelines) to the active command buffer until some specified limit is hit or the command buffer needs to be flushed for synchronization.
### Primitive Transforms[#](#primitive-transforms "Link to this heading")
Next, let’s add implementations for transformations in a `Primitive`. These transformations can be built on top of other operations, including the one we just defined:
/\*\* The Jacobian-vector product. \*/
std::vector Axpby::jvp(
const std::vector& primals,
const std::vector& tangents,
const std::vector& argnums) {
// Forward mode diff that pushes along the tangents
// The jvp transform on the primitive can built with ops
// that are scheduled on the same stream as the primitive
// If argnums = {0}, we only push along x in which case the
// jvp is just the tangent scaled by alpha
// Similarly, if argnums = {1}, the jvp is just the tangent
// scaled by beta
if (argnums.size() \> 1) {
auto scale \= argnums\[0\] \== 0 ? alpha\_ : beta\_;
auto scale\_arr \= array(scale, tangents\[0\].dtype());
return {multiply(scale\_arr, tangents\[0\], stream())};
}
// If, argnums = {0, 1}, we take contributions from both
// which gives us jvp = tangent\_x \* alpha + tangent\_y \* beta
else {
return {axpby(tangents\[0\], tangents\[1\], alpha\_, beta\_, stream())};
}
}
/\*\* The vector-Jacobian product. \*/
std::vector Axpby::vjp(
const std::vector& primals,
const std::vector& cotangents,
const std::vector& argnums,
const std::vector& /\* unused \*/) {
// Reverse mode diff
std::vector vjps;
for (auto arg : argnums) {
auto scale \= arg \== 0 ? alpha\_ : beta\_;
auto scale\_arr \= array(scale, cotangents\[0\].dtype());
vjps.push\_back(multiply(scale\_arr, cotangents\[0\], stream()));
}
return vjps;
}
Note, a transformation does not need to be fully defined to start using the `Primitive`.
/\*\* Vectorize primitive along given axis \*/
std::pair, std::vector> Axpby::vmap(
const std::vector& inputs,
const std::vector& axes) {
throw std::runtime\_error("\[Axpby\] vmap not implemented.");
}
Building and Binding[#](#building-and-binding "Link to this heading")
----------------------------------------------------------------------
Let’s look at the overall directory structure first.
extensions
├── axpby
│ ├── axpby.cpp
│ ├── axpby.h
│ └── axpby.metal
├── mlx\_sample\_extensions
│ └── \_\_init\_\_.py
├── bindings.cpp
├── CMakeLists.txt
└── setup.py
* `extensions/axpby/` defines the C++ extension library
* `extensions/mlx_sample_extensions` sets out the structure for the associated Python package
* `extensions/bindings.cpp` provides Python bindings for our operation
* `extensions/CMakeLists.txt` holds CMake rules to build the library and Python bindings
* `extensions/setup.py` holds the `setuptools` rules to build and install the Python package
### Binding to Python[#](#binding-to-python "Link to this heading")
We use [nanobind](https://nanobind.readthedocs.io/en/latest/)
to build a Python API for the C++ library. Since bindings for components such as [`mlx.core.array`](../python/_autosummary/mlx.core.array.html#mlx.core.array "mlx.core.array")
, [`mlx.core.stream`](../python/_autosummary/mlx.core.stream.html#mlx.core.stream "mlx.core.stream")
, etc. are already provided, adding our `axpby()` is simple.
NB\_MODULE(\_ext, m) {
m.doc() \= "Sample extension for MLX";
m.def(
"axpby",
&axpby,
"x"\_a,
"y"\_a,
"alpha"\_a,
"beta"\_a,
nb::kw\_only(),
"stream"\_a \= nb::none(),
R"(
Scale and sum two vectors element-wise
\`\`z = alpha \* x + beta \* y\`\`
Follows numpy style broadcasting between \`\`x\`\` and \`\`y\`\`
Inputs are upcasted to floats if needed
Args:
x (array): Input array.
y (array): Input array.
alpha (float): Scaling factor for \`\`x\`\`.
beta (float): Scaling factor for \`\`y\`\`.
Returns:
array: \`\`alpha \* x + beta \* y\`\`
)");
}
Most of the complexity in the above example comes from additional bells and whistles such as the literal names and doc-strings.
Warning
`mlx.core` must be imported before importing `mlx_sample_extensions` as defined by the nanobind module above to ensure that the casters for `mlx.core` components like [`mlx.core.array`](../python/_autosummary/mlx.core.array.html#mlx.core.array "mlx.core.array")
are available.
### Building with CMake[#](#building-with-cmake "Link to this heading")
Building the C++ extension library only requires that you `find_package(MLX CONFIG)` and then link it to your library.
\# Add library
add\_library(mlx\_ext)
\# Add sources
target\_sources(
mlx\_ext
PUBLIC
${CMAKE\_CURRENT\_LIST\_DIR}/axpby/axpby.cpp
)
\# Add include headers
target\_include\_directories(
mlx\_ext PUBLIC ${CMAKE\_CURRENT\_LIST\_DIR}
)
\# Link to mlx
target\_link\_libraries(mlx\_ext PUBLIC mlx)
We also need to build the attached Metal library. For convenience, we provide a `mlx_build_metallib()` function that builds a `.metallib` target given sources, headers, destinations, etc. (defined in `cmake/extension.cmake` and automatically imported with MLX package).
Here is what that looks like in practice:
\# Build metallib
if(MLX\_BUILD\_METAL)
mlx\_build\_metallib(
TARGET mlx\_ext\_metallib
TITLE mlx\_ext
SOURCES ${CMAKE\_CURRENT\_LIST\_DIR}/axpby/axpby.metal
INCLUDE\_DIRS ${PROJECT\_SOURCE\_DIR} ${MLX\_INCLUDE\_DIRS}
OUTPUT\_DIRECTORY ${CMAKE\_LIBRARY\_OUTPUT\_DIRECTORY}
)
add\_dependencies(
mlx\_ext
mlx\_ext\_metallib
)
endif()
Finally, we build the [nanobind](https://nanobind.readthedocs.io/en/latest/)
bindings
nanobind\_add\_module(
\_ext
NB\_STATIC STABLE\_ABI LTO NOMINSIZE
NB\_DOMAIN mlx
${CMAKE\_CURRENT\_LIST\_DIR}/bindings.cpp
)
target\_link\_libraries(\_ext PRIVATE mlx\_ext)
if(BUILD\_SHARED\_LIBS)
target\_link\_options(\_ext PRIVATE \-Wl,-rpath,@loader\_path)
endif()
### Building with `setuptools`[#](#building-with-setuptools "Link to this heading")
Once we have set out the CMake build rules as described above, we can use the build utilities defined in `mlx.extension`:
from mlx import extension
from setuptools import setup
if \_\_name\_\_ \== "\_\_main\_\_":
setup(
name\="mlx\_sample\_extensions",
version\="0.0.0",
description\="Sample C++ and Metal extensions for MLX primitives.",
ext\_modules\=\[extension.CMakeExtension("mlx\_sample\_extensions.\_ext")\],
cmdclass\={"build\_ext": extension.CMakeBuild},
packages\=\["mlx\_sample\_extensions"\],
package\_data\={"mlx\_sample\_extensions": \["\*.so", "\*.dylib", "\*.metallib"\]},
extras\_require\={"dev":\[\]},
zip\_safe\=False,
python\_requires\=">=3.8",
)
Note
We treat `extensions/mlx_sample_extensions` as the package directory even though it only contains a `__init__.py` to ensure the following:
* `mlx.core` must be imported before importing `_ext`
* The C++ extension library and the metal library are co-located with the python bindings and copied together if the package is installed
To build the package, first install the build dependencies with `pip install -r requirements.txt`. You can then build inplace for development using `python setup.py build_ext -j8 --inplace` (in `extensions/`)
This results in the directory structure:
extensions
├── mlx\_sample\_extensions
│ ├── \_\_init\_\_.py
│ ├── libmlx\_ext.dylib # C++ extension library
│ ├── mlx\_ext.metallib # Metal library
│ └── \_ext.cpython-3x-darwin.so # Python Binding
…
When you try to install using the command `python -m pip install .` (in `extensions/`), the package will be installed with the same structure as `extensions/mlx_sample_extensions` and the C++ and Metal library will be copied along with the Python binding since they are specified as `package_data`.
Usage[#](#usage "Link to this heading")
----------------------------------------
After installing the extension as described above, you should be able to simply import the Python package and play with it as you would any other MLX operation.
Let’s look at a simple script and its results:
import mlx.core as mx
from mlx\_sample\_extensions import axpby
a \= mx.ones((3, 4))
b \= mx.ones((3, 4))
c \= axpby(a, b, 4.0, 2.0, stream\=mx.cpu)
print(f"c shape: {c.shape}")
print(f"c dtype: {c.dtype}")
print(f"c correct: {mx.all(c \== 6.0).item()}")
Output:
c shape: \[3, 4\]
c dtype: float32
c correctness: True
### Results[#](#results "Link to this heading")
Let’s run a quick benchmark and see how our new `axpby` operation compares with the naive `simple_axpby()` we first defined.
import mlx.core as mx
from mlx\_sample\_extensions import axpby
import time
def simple\_axpby(x: mx.array, y: mx.array, alpha: float, beta: float) \-> mx.array:
return alpha \* x + beta \* y
M \= 4096
N \= 4096
x \= mx.random.normal((M, N))
y \= mx.random.normal((M, N))
alpha \= 4.0
beta \= 2.0
mx.eval(x, y)
def bench(f):
\# Warm up
for i in range(5):
z \= f(x, y, alpha, beta)
mx.eval(z)
\# Timed run
s \= time.time()
for i in range(100):
z \= f(x, y, alpha, beta)
mx.eval(z)
e \= time.time()
return 1000 \* (e \- s) / 100
simple\_time \= bench(simple\_axpby)
custom\_time \= bench(axpby)
print(f"Simple axpby: {simple\_time:.3f} ms | Custom axpby: {custom\_time:.3f} ms")
The results are `Simple axpby: 1.559 ms | Custom axpby: 0.774 ms`. We see modest improvements right away!
This operation is now good to be used to build other operations, in [`mlx.nn.Module`](../python/nn/module.html#mlx.nn.Module "mlx.nn.Module")
calls, and also as a part of graph transformations like `grad()`.
Scripts[#](#scripts "Link to this heading")
--------------------------------------------
Download the code
The full example code is available in [mlx](https://github.com/ml-explore/mlx/tree/main/examples/extensions/)
.
Contents
---
# Using MLX in C++ — MLX 0.24.0 documentation
[Skip to main content](#main-content)
Back to top
Ctrl+K
[](../index.html)
[](https://github.com/ml-explore/mlx "Source repository")
* [.rst](../_sources/dev/mlx_in_cpp.rst "Download source file")
* .pdf
Using MLX in C++
================
Using MLX in C++[#](#using-mlx-in-c "Link to this heading")
============================================================
You can use MLX in a C++ project with CMake.
Note
This guide is based one the following [example using MLX in C++](https://github.com/ml-explore/mlx/tree/main/examples/cmake_project)
First install MLX:
pip install \-U mlx
You can also install the MLX Python package from source or just the C++ library. For more information see the [documentation on installing MLX](../install.html#build-and-install)
.
Next make an example program in `example.cpp`:
#include
#include "mlx/mlx.h"
namespace mx \= mlx::core;
int main() {
auto x \= mx::array({1, 2, 3});
auto y \= mx::array({1, 2, 3});
std::cout << x + y << std::endl;
return 0;
}
The next step is to setup a CMake file in `CMakeLists.txt`:
cmake\_minimum\_required(VERSION 3.27)
project(example LANGUAGES CXX)
set(CMAKE\_CXX\_STANDARD 17)
set(CMAKE\_CXX\_STANDARD\_REQUIRED ON)
Depending on how you installed MLX, you may need to tell CMake where to find it.
If you installed MLX with Python, then add the following to the CMake file:
find\_package(
Python 3.9
COMPONENTS Interpreter Development.Module
REQUIRED)
execute\_process(
COMMAND "${Python\_EXECUTABLE}" \-m mlx \--cmake-dir
OUTPUT\_STRIP\_TRAILING\_WHITESPACE
OUTPUT\_VARIABLE MLX\_ROOT)
If you installed the MLX C++ package to a system path, then CMake should be able to find it. If you installed it to a non-standard location or CMake can’t find MLX then set `MLX_ROOT` to the location where MLX is installed:
set(MLX\_ROOT "/path/to/mlx/")
Next, instruct CMake to find MLX:
find\_package(MLX CONFIG REQUIRED)
Finally, add the `example.cpp` program as an executable and link MLX.
add\_executable(example example.cpp)
target\_link\_libraries(example PRIVATE mlx)
You can build the example with:
cmake \-B build \-DCMAKE\_BUILD\_TYPE\=Release
cmake \--build build
And run it with:
./build/example
Note `find_package(MLX CONFIG REQUIRED)` sets the following variables:
| | |
| --- | --- |Package Variables[#](#id1 "Link to this table")
| Variable | Description |
| --- | --- |
| MLX\_FOUND | `True` if MLX is found |
| MLX\_INCLUDE\_DIRS | Include directory |
| MLX\_LIBRARIES | Libraries to link against |
| MLX\_CXX\_FLAGS | Additional compiler flags |
| MLX\_BUILD\_ACCELERATE | `True` if MLX was built with Accelerate |
| MLX\_BUILD\_METAL | `True` if MLX was built with Metal |
---
# Unknown
MLX === MLX is a NumPy-like array framework designed for efficient and flexible machine learning on Apple silicon, brought to you by Apple machine learning research. The Python API closely follows NumPy with a few exceptions. MLX also has a fully featured C++ API which closely follows the Python API. The main differences between MLX and NumPy are: - \*\*Composable function transformations\*\*: MLX has composable function transformations for automatic differentiation, automatic vectorization, and computation graph optimization. - \*\*Lazy computation\*\*: Computations in MLX are lazy. Arrays are only materialized when needed. - \*\*Multi-device\*\*: Operations can run on any of the supported devices (CPU, GPU, ...) The design of MLX is inspired by frameworks like \`PyTorch \`\_, \`Jax \`\_, and \`ArrayFire \`\_. A notable difference from these frameworks and MLX is the \*unified memory model\*. Arrays in MLX live in shared memory. Operations on MLX arrays can be performed on any of the supported device types without performing data copies. Currently supported device types are the CPU and GPU. .. toctree:: :caption: Install :maxdepth: 1 install .. toctree:: :caption: Usage :maxdepth: 1 usage/quick\_start usage/lazy\_evaluation usage/unified\_memory usage/indexing usage/saving\_and\_loading usage/function\_transforms usage/compile usage/numpy usage/distributed usage/using\_streams usage/export .. toctree:: :caption: Examples :maxdepth: 1 examples/linear\_regression examples/mlp examples/llama-inference .. toctree:: :caption: Python API Reference :maxdepth: 1 python/array python/data\_types python/devices\_and\_streams python/export python/ops python/random python/transforms python/fast python/fft python/linalg python/metal python/nn python/optimizers python/distributed python/tree\_utils .. toctree:: :caption: C++ API Reference :maxdepth: 1 cpp/ops .. toctree:: :caption: Further Reading :maxdepth: 1 dev/extensions dev/metal\_debugger dev/custom\_metal\_kernels dev/mlx\_in\_cpp
---
# Unknown
.. \_build\_and\_install: Build and Install ================= Python Installation ------------------- MLX is available on PyPI. All you have to do to use MLX with your own Apple silicon computer is .. code-block:: shell pip install mlx To install from PyPI you must meet the following requirements: - Using an M series chip (Apple silicon) - Using a native Python >= 3.9 - macOS >= 13.5 .. note:: MLX is only available on devices running macOS >= 13.5 It is highly recommended to use macOS 14 (Sonoma) MLX is also available on conda-forge. To install MLX with conda do: .. code-block:: shell conda install conda-forge::mlx Troubleshooting ^^^^^^^^^^^^^^^ \*My OS and Python versions are in the required range but pip still does not find a matching distribution.\* Probably you are using a non-native Python. The output of .. code-block:: shell python -c "import platform; print(platform.processor())" should be \`\`arm\`\`. If it is \`\`i386\`\` (and you have M series machine) then you are using a non-native Python. Switch your Python to a native Python. A good way to do this is with \`Conda \`\_. Build from source ----------------- Build Requirements ^^^^^^^^^^^^^^^^^^ - A C++ compiler with C++17 support (e.g. Clang >= 5.0) - \`cmake \`\_ -- version 3.25 or later, and \`\`make\`\` - Xcode >= 15.0 and macOS SDK >= 14.0 .. note:: Ensure your shell environment is native \`\`arm\`\`, not \`\`x86\`\` via Rosetta. If the output of \`\`uname -p\`\` is \`\`x86\`\`, see the :ref:\`troubleshooting section \` below. Python API ^^^^^^^^^^ To build and install the MLX python library from source, first, clone MLX from \`its GitHub repo \`\_: .. code-block:: shell git clone git@github.com:ml-explore/mlx.git mlx && cd mlx Then simply build and install MLX using pip: .. code-block:: shell CMAKE\_BUILD\_PARALLEL\_LEVEL=8 pip install . For developing, install the package with development dependencies, and use an editable install: .. code-block:: shell CMAKE\_BUILD\_PARALLEL\_LEVEL=8 pip install -e ".\[dev\]" Once the development dependencies are installed, you can build faster with: .. code-block:: shell CMAKE\_BUILD\_PARALLEL\_LEVEL=8 python setup.py build\_ext --inplace Run the tests with: .. code-block:: shell python -m unittest discover python/tests Optional: Install stubs to enable auto completions and type checking from your IDE: .. code-block:: shell python setup.py generate\_stubs C++ API ^^^^^^^ Currently, MLX must be built and installed from source. Similarly to the python library, to build and install the MLX C++ library start by cloning MLX from \`its GitHub repo \`\_: .. code-block:: shell git clone git@github.com:ml-explore/mlx.git mlx && cd mlx Create a build directory and run CMake and make: .. code-block:: shell mkdir -p build && cd build cmake .. && make -j Run tests with: .. code-block:: shell make test Install with: .. code-block:: shell make install Note that the built \`\`mlx.metallib\`\` file should be either at the same directory as the executable statically linked to \`\`libmlx.a\`\` or the preprocessor constant \`\`METAL\_PATH\`\` should be defined at build time and it should point to the path to the built metal library. .. list-table:: Build Options :widths: 25 8 :header-rows: 1 \* - Option - Default \* - MLX\_BUILD\_TESTS - ON \* - MLX\_BUILD\_EXAMPLES - OFF \* - MLX\_BUILD\_BENCHMARKS - OFF \* - MLX\_BUILD\_METAL - ON \* - MLX\_BUILD\_CPU - ON \* - MLX\_BUILD\_PYTHON\_BINDINGS - OFF \* - MLX\_METAL\_DEBUG - OFF \* - MLX\_BUILD\_SAFETENSORS - ON \* - MLX\_BUILD\_GGUF - ON \* - MLX\_METAL\_JIT - OFF .. note:: If you have multiple Xcode installations and wish to use a specific one while building, you can do so by adding the following environment variable before building .. code-block:: shell export DEVELOPER\_DIR="/path/to/Xcode.app/Contents/Developer/" Further, you can use the following command to find out which macOS SDK will be used .. code-block:: shell xcrun -sdk macosx --show-sdk-version Binary Size Minimization ~~~~~~~~~~~~~~~~~~~~~~~~ To produce a smaller binary use the CMake flags \`\`CMAKE\_BUILD\_TYPE=MinSizeRel\`\` and \`\`BUILD\_SHARED\_LIBS=ON\`\`. The MLX CMake build has several additional options to make smaller binaries. For example, if you don't need the CPU backend or support for safetensors and GGUF, you can do: .. code-block:: shell cmake .. \\ -DCMAKE\_BUILD\_TYPE=MinSizeRel \\ -DBUILD\_SHARED\_LIBS=ON \\ -DMLX\_BUILD\_CPU=OFF \\ -DMLX\_BUILD\_SAFETENSORS=OFF \\ -DMLX\_BUILD\_GGUF=OFF \\ -DMLX\_METAL\_JIT=ON THE \`\`MLX\_METAL\_JIT\`\` flag minimizes the size of the MLX Metal library which contains pre-built GPU kernels. This substantially reduces the size of the Metal library by run-time compiling kernels the first time they are used in MLX on a given machine. Note run-time compilation incurs a cold-start cost which can be anwywhere from a few hundred millisecond to a few seconds depending on the application. Once a kernel is compiled, it will be cached by the system. The Metal kernel cache persists across reboots. Troubleshooting ^^^^^^^^^^^^^^^ Metal not found ~~~~~~~~~~~~~~~ You see the following error when you try to build: .. code-block:: shell error: unable to find utility "metal", not a developer tool or in PATH To fix this, first make sure you have Xcode installed: .. code-block:: shell xcode-select --install Then set the active developer directory: .. code-block:: shell sudo xcode-select --switch /Applications/Xcode.app/Contents/Developer x86 Shell ~~~~~~~~~ .. \_build shell: If the output of \`\`uname -p\`\` is \`\`x86\`\` then your shell is running as x86 via Rosetta instead of natively. To fix this, find the application in Finder (\`\`/Applications\`\` for iTerm, \`\`/Applications/Utilities\`\` for Terminal), right-click, and click “Get Info”. Uncheck “Open using Rosetta”, close the “Get Info” window, and restart your terminal. Verify the terminal is now running natively the following command: .. code-block:: shell $ uname -p arm Also check that cmake is using the correct architecture: .. code-block:: shell $ cmake --system-information | grep CMAKE\_HOST\_SYSTEM\_PROCESSOR CMAKE\_HOST\_SYSTEM\_PROCESSOR "arm64" If you see \`\`"x86\_64"\`\`, try re-installing \`\`cmake\`\`. If you see \`\`"arm64"\`\` but the build errors out with "Building for x86\_64 on macOS is not supported." wipe your build cache with \`\`rm -rf build/\`\` and try again.
---
# Launching Distributed Programs — MLX 0.24.0 documentation
[Skip to main content](#main-content)
Back to top
Ctrl+K
[](../index.html)
[](https://github.com/ml-explore/mlx "Source repository")
* [.rst](../_sources/usage/launching_distributed.rst "Download source file")
* .pdf
Launching Distributed Programs
==============================
Contents
--------
Launching Distributed Programs[#](#launching-distributed-programs "Link to this heading")
==========================================================================================
Installing the MLX python package provides a helper script `mlx.launch` that can be used to run python scripts distributed on several nodes. It allows launching using either the MPI backend or the ring backend. See the [distributed docs](distributed.html)
for the different backends.
Usage[#](#usage "Link to this heading")
----------------------------------------
The minimal usage example of `mlx.launch` is simply
mlx.launch \--hosts ip1,ip2 my\_script.py
or for testing on localhost
mlx.launch \-n 2 my\_script.py
The `mlx.launch` command connects to the provided host and launches the input script on each host. It monitors each of the launched processes and terminates the rest if one of them fails unexpectedly or if `mlx.launch` is terminated. It also takes care of forwarding the output of each remote process to stdout and stderr respectively.
### Providing Hosts[#](#providing-hosts "Link to this heading")
Hosts can be provided as command line arguments, like above, but the way that allows to fully define a list of hosts is via a JSON hostfile. The hostfile has a very simple schema. It is simply a list of objects that define each host via a hostname to ssh to and a list of IPs to utilize for the communication.
\[\
{"ssh": "hostname1", "ips": \["123.123.1.1", "123.123.2.1"\]},\
{"ssh": "hostname2", "ips": \["123.123.1.2", "123.123.2.2"\]}\
\]
You can use `mlx.distributed_config --over ethernet` to create a hostfile with IPs corresponding to the `en0` interface.
### Setting up Remote Hosts[#](#setting-up-remote-hosts "Link to this heading")
In order to be able to launch the script on each host we need to be able to connect via ssh. Moreover the input script and python binary need to be on each host and on the same path. A good checklist to debug errors is the following:
* `ssh hostname` works without asking for password or host confirmation
* the python binary is available on all hosts at the same path. You can use `mlx.launch --print-python` to see what that path is.
* the script you want to run is available on all hosts at the same path
MPI Specifics[#](#mpi-specifics "Link to this heading")
--------------------------------------------------------
One can use MPI by passing `--backend mpi` to `mlx.launch`. In that case, `mlx.launch` is a thin wrapper over `mpirun`. Moreover,
* The IPs in the hostfile are ignored
* The ssh connectivity requirement is stronger as every node needs to be able to connect to every other node
* `mpirun` needs to be available on every node at the same path
Finally, one can pass arguments to `mpirun` using `--mpi-arg`. For instance to choose a specific interface for the byte-transfer-layer of MPI we can call `mlx.launch` as follows:
mlx.launch \--backend mpi \--mpi-arg '--mca btl\_tcp\_if\_include en0' \--hostfile hosts.json my\_script.py
Ring Specifics[#](#ring-specifics "Link to this heading")
----------------------------------------------------------
The ring backend, which is also the default backend, can be explicitly selected with the argument `--backend ring`. The ring backend has some specific requirements and arguments that are different to MPI:
* The argument `--hosts` only accepts IPs and not hostnames. If we need to ssh to a hostname that does not correspond to the IP we want to bind to we have to provide a hostfile.
* `--starting-port` defines the port to bind to on the remote hosts. Specifically rank 0 for the first IP will use this port and each subsequent IP or rank will add 1 to this port.
* `--connections-per-ip` allows us to increase the number of connections between neighboring nodes. This corresponds to `--mca btl_tcp_links 2` for `mpirun`.
Contents
---