### Install MLX from Source with Development Dependencies

Source: https://ml-explore.github.io/mlx/build/html/install.html

Install MLX from source with development dependencies for an editable install. This is recommended for active development.

```bash
pip install -e ".[dev]"
```

--------------------------------

### Build and Install MLX from Source (Python)

Source: https://ml-explore.github.io/mlx/build/html/install.html

Build and install the MLX Python library from source using pip. This is the standard method after cloning the repository.

```bash
pip install .
```

--------------------------------

### Example C++ Program with MLX

Source: https://ml-explore.github.io/mlx/build/html/dev/mlx_in_cpp.html

A basic C++ program demonstrating the use of MLX arrays and operations. Ensure MLX is installed and accessible.

```cpp
#include <iostream>

#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;
}

```

--------------------------------

### Install BLAS and LAPACK Headers on Ubuntu for CPU-only Build

Source: https://ml-explore.github.io/mlx/build/html/install.html

Install the necessary BLAS, LAPACK, and LAPACKE development headers on Ubuntu for building MLX with CPU support only.

```bash
apt-get update -y
apt-get install libblas-dev liblapack-dev liblapacke-dev -y

```

--------------------------------

### Ring Backend Hostfile Configuration

Source: https://ml-explore.github.io/mlx/build/html/usage/distributed.html

Example JSON structure for the `MLX_HOSTFILE` environment variable, specifying IPs and ports for each rank in a distributed setup.

```json
[
  ["123.123.1.1:5000", "123.123.1.2:5000"],
  ["123.123.2.1:5000", "123.123.2.2:5000"],
  ["123.123.3.1:5000", "123.123.3.2:5000"],
  ["123.123.4.1:5000", "123.123.4.2:5000"]
]
```

--------------------------------

### MaxPool1d Example

Source: https://ml-explore.github.io/mlx/build/html/python/nn/_autosummary/mlx.nn.MaxPool1d.html

Demonstrates how to initialize and use the MaxPool1d layer for 1D max pooling.

```python
>>> import mlx.core as mx
>>> import mlx.nn.layers as nn
>>> x = mx.random.normal(shape=(4, 16, 5))
>>> pool = nn.MaxPool1d(kernel_size=2, stride=2)
>>> pool(x)
```

--------------------------------

### Install MLX C++ Library

Source: https://ml-explore.github.io/mlx/build/html/install.html

Install the MLX C++ library after building. This command makes the library available for linking into other projects.

```bash
make install
```

--------------------------------

### start_capture()

Source: https://ml-explore.github.io/mlx/build/html/python/_autosummary/mlx.core.metal.start_capture.html

Starts a Metal capture and saves the trace to the specified path.

```APIDOC
## start_capture()

### Description
Starts a Metal capture. The capture will be saved to the provided path.

### Parameters
#### Path Parameters
- **path** (str) - Required - The path to save the capture which should have the extension `.gputrace`.

### Returns
None
```

--------------------------------

### Run MLX C++ Example

Source: https://ml-explore.github.io/mlx/build/html/dev/mlx_in_cpp.html

Command to execute the compiled C++ program. Assumes the build output is in the 'build' directory.

```bash
./build/example

```

--------------------------------

### start_capture

Source: https://ml-explore.github.io/mlx/build/html/python/metal.html

Starts a Metal capture session to a specified file path.

```APIDOC
## start_capture(path)

### Description
Start a Metal capture to the specified file path.

### Method
N/A (Function call)

### Parameters
#### Path Parameters
- **path** (string) - Required - The file path where the capture will be saved.
```

--------------------------------

### Install Xcode Command Line Tools

Source: https://ml-explore.github.io/mlx/build/html/install.html

Install the Xcode command line tools, which are required for Metal development and may be missing if Xcode is not fully installed.

```bash
xcode-select --install

```

--------------------------------

### Example Usage of Module

Source: https://ml-explore.github.io/mlx/build/html/python/nn/module.html

Demonstrates how to define a simple MLP using the Module base class and manage its parameters.

```python
import mlx.core as mx
import mlx.nn as nn

class MyMLP(nn.Module):
    def __init__(self, in_dims: int, out_dims: int, hidden_dims: int = 16):
        super().__init__()

        self.in_proj = nn.Linear(in_dims, hidden_dims)
        self.out_proj = nn.Linear(hidden_dims, out_dims)

    def __call__(self, x):
        x = self.in_proj(x)
        x = mx.maximum(x, 0)
        return self.out_proj(x)

model = MyMLP(2, 1)

# Evaluate and initialize parameters
mx.eval(model.parameters())

# Modify a parameter
model.in_proj.weight = model.in_proj.weight * 2
mx.eval(model.parameters())
```

--------------------------------

### JACCL Backend Ibverbs Devices Configuration

Source: https://ml-explore.github.io/mlx/build/html/usage/distributed.html

Example JSON structure for the `MLX_IBV_DEVICES` environment variable, specifying the ibverbs device names for inter-node connections in a JACCL distributed setup.

```json
[
   [null, "rdma_en5", "rdma_en4", "rdma_en3"],
   ["rdma_en5", null, "rdma_en3", "rdma_en4"],
   ["rdma_en4", "rdma_en3", null, "rdma_en5"],
   ["rdma_en3", "rdma_en4", "rdma_en5", null]
]
```

--------------------------------

### Backend Initialization Logic Examples

Source: https://ml-explore.github.io/mlx/build/html/usage/distributed.html

Demonstrates different scenarios for initializing communication backends in MLX. Subsequent calls to `init()` without arguments will return the already initialized backend.

```python
# 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!
```

--------------------------------

### Install MLX with Pip

Source: https://ml-explore.github.io/mlx/build/html/install.html

Install the MLX Python package from PyPI. Requires Apple silicon, Python >= 3.10, and macOS >= 14.0.

```bash
pip install mlx
```

--------------------------------

### Install CUDA Toolkit and Dependencies on Ubuntu

Source: https://ml-explore.github.io/mlx/build/html/install.html

Install the CUDA toolkit, BLAS, LAPACK, LAPACKE, and cuDNN development libraries on Ubuntu for building MLX with CUDA support.

```bash
wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.1-1_all.deb
dpkg -i cuda-keyring_1.1-1_all.deb
apt-get update -y
apt-get -y install cuda-toolkit-12-9
apt-get install libblas-dev liblapack-dev liblapacke-dev libcudnn9-dev-cuda-12 -y

```

--------------------------------

### 3D Average Pooling Example

Source: https://ml-explore.github.io/mlx/build/html/python/nn/_autosummary/mlx.nn.AvgPool3d.html

Demonstrates how to initialize and use the AvgPool3d layer for 3D average pooling. Ensure mlx.core and mlx.nn.layers are imported.

```python
>>> import mlx.core as mx
>>> import mlx.nn.layers as nn
>>> x = mx.random.normal(shape=(8, 16, 32, 32, 4))
>>> pool = nn.AvgPool3d(kernel_size=2, stride=2)
>>> pool(x)
```

--------------------------------

### Install MLX Python Package

Source: https://ml-explore.github.io/mlx/build/html/dev/mlx_in_cpp.html

Use pip to install the MLX Python package. This is a prerequisite for using MLX in C++ if you are installing from Python.

```bash
pip install -U mlx

```

--------------------------------

### CMakeLists.txt Setup for C++ Project

Source: https://ml-explore.github.io/mlx/build/html/dev/mlx_in_cpp.html

Initial CMake configuration for a C++ project. Sets the minimum CMake version and project name, and specifies C++20 standard.

```cmake
cmake_minimum_required(VERSION 3.27)

project(example LANGUAGES CXX)

set(CMAKE_CXX_STANDARD 20)
set(CMAKE_CXX_STANDARD_REQUIRED ON)

```

--------------------------------

### Install MLX CPU-only (Linux)

Source: https://ml-explore.github.io/mlx/build/html/install.html

Install a CPU-only version of the MLX Python package for Linux. Requires a compatible Linux distribution and Python version.

```bash
pip install mlx[cpu]
```

--------------------------------

### Install MLX with CUDA Support

Source: https://ml-explore.github.io/mlx/build/html/install.html

Install the MLX Python package with CUDA support from PyPI. Requires specific Nvidia hardware, drivers, CUDA toolkit, Linux distribution, and Python version.

```bash
pip install mlx[cuda12]
```

--------------------------------

### InstanceNorm Usage Example

Source: https://ml-explore.github.io/mlx/build/html/python/nn/_autosummary/mlx.nn.InstanceNorm.html

Demonstrates how to initialize and use the InstanceNorm layer with a random input tensor. Ensure mlx.core and mlx.nn are imported.

```python
>>> import mlx.core as mx
>>> import mlx.nn as nn
>>> x = mx.random.normal((8, 4, 4, 16))
>>> inorm = nn.InstanceNorm(dims=16)
>>> output = inorm(x)
```

--------------------------------

### Launch Script on Localhost

Source: https://ml-explore.github.io/mlx/build/html/usage/launching_distributed.html

Launch a script on the local machine using multiple processes with `mlx.launch`. This is useful for testing or multi-GPU setups on a single node.

```bash
mlx.launch -n 2 my_script.py
```

--------------------------------

### fsdp_apply_gradients() Example

Source: https://ml-explore.github.io/mlx/build/html/python/_autosummary/mlx.nn.fsdp_apply_gradients.html

Demonstrates how to use `fsdp_apply_gradients` with and without gradient clipping. Ensure `grads`, `params`, and `optimizer` are properly initialized. The `model.update()` call is a placeholder for applying the updated parameters.

```python
optimizer = optim.SGD(learning_rate=0.01)
>>> # Without gradient clipping
>>> updated_params = fsdp_apply_gradients(grads, params, optimizer)
>>> model.update(updated_params)
>>>
>>> # With gradient clipping
>>> updated_params, grad_norm = fsdp_apply_gradients(
...     grads, params, optimizer, max_norm=1.0
... )
>>> model.update(updated_params)
```

--------------------------------

### Setuptools Configuration for MLX Extensions

Source: https://ml-explore.github.io/mlx/build/html/dev/extensions.html

Configures the Python package setup using setuptools to include CMake-built extensions. Specifies package name, version, description, and extension modules.

```python
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",
    )

```

--------------------------------

### Visualize Thunderbolt Connectivity

Source: https://ml-explore.github.io/mlx/build/html/usage/distributed.html

Uses `mlx.distributed_config` to visualize the Thunderbolt network topology between specified hosts. Requires `dot` and `Preview` to be installed.

```bash
mlx.distributed_config --verbose \
     --hosts m3-ultra-1,m3-ultra-2,m3-ultra-3,m3-ultra-4 \
     --over thunderbolt --dot | dot -Tpng | open -f -a Preview

```

--------------------------------

### Import MLX Core and Setup Metadata

Source: https://ml-explore.github.io/mlx/build/html/examples/linear_regression.html

Imports the core MLX package and sets up metadata for the linear regression problem, including dataset dimensions and optimization parameters.

```python
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
```

--------------------------------

### Deferred Computation Example

Source: https://ml-explore.github.io/mlx/build/html/usage/lazy_evaluation.html

Demonstrates how MLX defers computation. Only `a` is computed, while `expensive_fun(a)` is not executed because its output is ignored.

```python
def fun(x):
    a = fun1(x)
    b = expensive_fun(a)
    return a, b

y, _ = fun(x)

```

--------------------------------

### Verify Thunderbolt RDMA Devices

Source: https://ml-explore.github.io/mlx/build/html/usage/distributed.html

After enabling RDMA, use `ibv_devices` to list the available RDMA network interfaces. This command should output device names and their GUIDs.

```bash
~ % ibv_devices
device                 node GUID
------              ----------------
rdma_en2            8096a9d9edbaac05
rdma_en3            8196a9d9edbaac05
rdma_en5            8396a9d9edbaac05
rdma_en4            8296a9d9edbaac05
rdma_en6            8496a9d9edbaac05
rdma_en7            8596a9d9edbaac05

```

--------------------------------

### Setup Problem Parameters and Load Data

Source: https://ml-explore.github.io/mlx/build/html/examples/mlp.html

Sets hyperparameters for the MLP model and loads the MNIST dataset using a custom mnist loader. Converts loaded data to MLX arrays.

```python
# Load the data
import mnist
train_images, train_labels, test_images, test_labels = map(
    mx.array, mnist.mnist()
)

```

--------------------------------

### Launch Distributed MLX Program with MPI

Source: https://ml-explore.github.io/mlx/build/html/usage/distributed.html

Use `mlx.launch` with the `mpi` backend to run distributed MLX programs. This command simplifies setup by managing paths and shared libraries.

```bash
$ mlx.launch --backend mpi -n 2 test.py
```

--------------------------------

### Standard In-place Add Operation Example

Source: https://ml-explore.github.io/mlx/build/html/python/_autosummary/mlx.core.array.at.html

Illustrates the behavior of standard in-place addition (+=) on an array with duplicate indices, showing that only the last update for a given index is effectively applied.

```python
>>> a = mx.array([0, 0])
>>> idx = mx.array([0, 1, 0, 1])
>>> a[idx] += 1
>>> a
array([1, 1], dtype=int32)
```

--------------------------------

### init

Source: https://ml-explore.github.io/mlx/build/html/python/optimizers/_autosummary/mlx.optimizers.MultiOptimizer.html

Initialize the optimizer's state.

```APIDOC
## init

### Description
Initialize the optimizer's state.

### Parameters
* **parameters** - The parameters to initialize the optimizer's state with.
```

--------------------------------

### Install OpenMPI with Conda

Source: https://ml-explore.github.io/mlx/build/html/usage/distributed.html

Install the OpenMPI package using the Conda package manager from the conda-forge channel.

```bash
$ conda install conda-forge::openmpi
```

--------------------------------

### Auto-configure Nodes and Generate Hostfile

Source: https://ml-explore.github.io/mlx/build/html/usage/distributed.html

Automatically sets up nodes for JACCL backend, checks for a valid mesh, and generates a hostfile for `mlx.launch`. Requires SSH access and potentially password-less sudo.

```bash
mlx.distributed_config --verbose \
     --hosts m3-ultra-1,m3-ultra-2,m3-ultra-3,m3-ultra-4 \
     --over thunderbolt --backend jaccl \
     --auto-setup --output m3-ultra-jaccl.json

```

--------------------------------

### Create and Apply Uniform Initializer

Source: https://ml-explore.github.io/mlx/build/html/python/nn/init.html

Demonstrates how to create a uniform initializer and apply it to create an initialized array. The initializer function takes an input array and returns an initialized output.

```python
import mlx.core as mx
import mlx.nn as nn

init_fn = nn.init.uniform()

# Produces a [2, 2] uniform matrix
param = init_fn(mx.zeros((2, 2)))
```

--------------------------------

### Set MLX_ROOT for Non-Standard Installation

Source: https://ml-explore.github.io/mlx/build/html/dev/mlx_in_cpp.html

Manually set the MLX_ROOT variable in CMake if MLX is installed in a non-standard location or cannot be found automatically.

```cmake
set(MLX_ROOT "/path/to/mlx/")

```

--------------------------------

### Scaled Dot Product Attention Example

Source: https://ml-explore.github.io/mlx/build/html/python/_autosummary/mlx.core.fast.scaled_dot_product_attention.html

Demonstrates the usage of scaled_dot_product_attention with causal masking. Ensure inputs q, k, and v have the correct dimensions and scale is calculated appropriately.

```python
B = 2
N_q = N_kv = 32
T_q = T_kv = 1000
D = 128

q = mx.random.normal(shape=(B, N_q, T_q, D))
k = mx.random.normal(shape=(B, N_kv, T_kv, D))
v = mx.random.normal(shape=(B, N_kv, T_kv, D))
scale = D ** -0.5
out = mx.fast.scaled_dot_product_attention(q, k, v, scale=scale, mask="causal")

```

--------------------------------

### Compiled Training Loop with State Capture

Source: https://ml-explore.github.io/mlx/build/html/usage/compile.html

This example shows how to compile a full training step, including forward pass, gradient calculation, and optimizer update, by capturing model and optimizer states as both inputs and outputs.

```python
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)
```

--------------------------------

### array slice(const array &a, const array &start, std::vector<int> axes, Shape slice_size, StreamOrDevice s = {})

Source: https://ml-explore.github.io/mlx/build/html/cpp/ops.html

Slice an array with dynamic starting indices.

```APIDOC
## array slice(const array &a, const array &start, std::vector<int> axes, Shape slice_size, StreamOrDevice s = {})

### Description
Slice an array with dynamic starting indices.

### Parameters
- **a** (const array &) - The input array.
- **start** (const array &) - The dynamic starting indices for slicing.
- **axes** (std::vector<int>) - The axes along which to slice.
- **slice_size** (Shape) - The size of the slice.
- **s** (StreamOrDevice) - Optional stream or device for computation.
```

--------------------------------

### array slice(const array &a, std::initializer_list<int> start, Shape stop, Shape strides, StreamOrDevice s = {})

Source: https://ml-explore.github.io/mlx/build/html/cpp/ops.html

Slice an array with specified start, stop, and strides.

```APIDOC
## array slice(const array &a, std::initializer_list<int> start, Shape stop, Shape strides, StreamOrDevice s = {})

### Description
Slice an array with specified start, stop, and strides.

### Parameters
- **a** (const array &) - The input array.
- **start** (std::initializer_list<int>) - The starting indices for slicing.
- **stop** (Shape) - The stopping indices for slicing.
- **strides** (Shape) - The strides for slicing.
- **s** (StreamOrDevice) - Optional stream or device for computation.
```

--------------------------------

### array slice_update(const array &src, const array &update, const array &start, std::vector<int> axes, StreamOrDevice s = {})

Source: https://ml-explore.github.io/mlx/build/html/cpp/ops.html

Update a slice from the source array with dynamic starting indices.

```APIDOC
## array slice_update(const array &src, const array &update, const array &start, std::vector<int> axes, StreamOrDevice s = {})

### Description
Update a slice from the source array with dynamic starting indices.

### Parameters
- **src** (const array &) - The source array.
- **update** (const array &) - The array containing the update values.
- **start** (const array &) - The dynamic starting indices for the slice to update.
- **axes** (std::vector<int>) - The axes along which to update.
- **s** (StreamOrDevice) - Optional stream or device for computation.
```

--------------------------------

### array slice_update(const array &src, const array &update, Shape start, Shape stop, Shape strides, StreamOrDevice s = {})

Source: https://ml-explore.github.io/mlx/build/html/cpp/ops.html

Update a slice from the source array with specified start, stop, and strides.

```APIDOC
## array slice_update(const array &src, const array &update, Shape start, Shape stop, Shape strides, StreamOrDevice s = {})

### Description
Update a slice from the source array with specified start, stop, and strides.

### Parameters
- **src** (const array &) - The source array.
- **update** (const array &) - The array containing the update values.
- **start** (Shape) - The starting indices for the slice to update.
- **stop** (Shape) - The stopping indices for the slice to update.
- **strides** (Shape) - The strides for the slice update.
- **s** (StreamOrDevice) - Optional stream or device for computation.
```

--------------------------------

### Initialize and Use AvgPool2d

Source: https://ml-explore.github.io/mlx/build/html/python/nn/_autosummary/mlx.nn.AvgPool2d.html

Demonstrates how to initialize an AvgPool2d layer with a kernel size and stride, and then apply it to an input tensor. Ensure mlx.core and mlx.nn.layers are imported.

```python
import mlx.core as mx
import mlx.nn.layers as nn
x = mx.random.normal(shape=(8, 32, 32, 4))
pool = nn.AvgPool2d(kernel_size=2, stride=2)
pool(x)
```

--------------------------------

### default_device()

Source: https://ml-explore.github.io/mlx/build/html/python/_autosummary/mlx.core.default_device.html

Get the default device.

```APIDOC
## default_device()

### Description
Get the default device.

### Signature
default_device() -> Device
```

--------------------------------

### arange

Source: https://ml-explore.github.io/mlx/build/html/cpp/ops.html

Creates a 1D array of numbers starting at `start`, stopping at `stop`, and stepping by `step`. Various overloads exist to handle optional parameters like `dtype` and `StreamOrDevice`, and to accept integer or double types for arguments.

```APIDOC
## arange

### Description
A 1D array of numbers starting at `start` (optional), stopping at stop, stepping by `step` (optional).

### Signature
```cpp
array arange(double start, double stop, double step, Dtype dtype, StreamOrDevice s = {})
array arange(double start, double stop, double step, StreamOrDevice s = {})
array arange(double start, double stop, Dtype dtype, StreamOrDevice s = {})
array arange(double start, double stop, StreamOrDevice s = {})
array arange(double stop, Dtype dtype, StreamOrDevice s = {})
array arange(double stop, StreamOrDevice s = {})
array arange(int start, int stop, int step, StreamOrDevice s = {})
array arange(int start, int stop, StreamOrDevice s = {})
array arange(int stop, StreamOrDevice s = {})
```
```

--------------------------------

### Define Hosts with JSON Hostfile

Source: https://ml-explore.github.io/mlx/build/html/usage/launching_distributed.html

Specify hosts for `mlx.launch` using a JSON hostfile. Each entry defines the SSH target and available communication IPs for a node.

```json
[
    {"ssh": "hostname1", "ips": ["123.123.1.1", "123.123.2.1"]},
    {"ssh": "hostname2", "ips": ["123.123.1.2", "123.123.2.2"]}
]
```

--------------------------------

### default_stream()

Source: https://ml-explore.github.io/mlx/build/html/python/_autosummary/mlx.core.default_stream.html

Get the device’s default stream.

```APIDOC
## default_stream()

### Description
Get the device’s default stream.

### Signature
default_stream(_device : Device_) -> Stream
```

--------------------------------

### get_printoptions

Source: https://ml-explore.github.io/mlx/build/html/python/printoptions.html

Get global printing precision for array formatting.

```APIDOC
## get_printoptions

### Description
Retrieves the current global settings for array printing precision.

### Method
`get_printoptions()`

### Parameters
None

### Response
#### Success Response (200)
- **precision** (int) - The current global printing precision.
```

--------------------------------

### Initializers

Source: https://ml-explore.github.io/mlx/build/html/python/nn.html

Weight initialization functions in the mlx.nn.init module.

```APIDOC
## Initializers

This section lists the weight initialization functions available in `mlx.nn.init`.

### `mlx.nn.init.constant()`

Initializes weights with a constant value.

### `mlx.nn.init.normal()`

Initializes weights from a normal distribution.

### `mlx.nn.init.uniform()`

Initializes weights from a uniform distribution.

### `mlx.nn.init.identity()`

Initializes weights as an identity matrix.

### `mlx.nn.init.glorot_normal()`

Initializes weights using Glorot (Xavier) normal initialization.

### `mlx.nn.init.glorot_uniform()`

Initializes weights using Glorot (Xavier) uniform initialization.

### `mlx.nn.init.he_normal()`

Initializes weights using He normal initialization.

### `mlx.nn.init.he_uniform()`

Initializes weights using He uniform initialization.
```

--------------------------------

### Element Count

Source: https://ml-explore.github.io/mlx/build/html/cpp/ops.html

Function to get the total number of elements in an array.

```APIDOC
## Element Count

### Description
Returns the total number of elements in an array.

### Operations
- `number_of_elements()`: Returns the total number of elements in the array.
```

--------------------------------

### init

Source: https://ml-explore.github.io/mlx/build/html/python/distributed.html

Initializes the communication backend and creates the global communication group.

```APIDOC
## init([strict, backend])

### Description
Initializes the communication backend and creates the global communication group.

### Parameters
#### Query Parameters
- **strict** (bool) - Optional - If True, initialization will fail if the backend is not available.
- **backend** (string) - Optional - The communication backend to initialize (e.g., 'mpi').
```

--------------------------------

### mx.random.key

Source: https://ml-explore.github.io/mlx/build/html/python/random.html

Get a PRNG key from a seed. Takes a `seed` argument.

```APIDOC
## mx.random.key

### Description
Get a PRNG key from a seed.

### Parameters
- **seed** (int) - The seed for the PRNG key.
```

--------------------------------

### init()

Source: https://ml-explore.github.io/mlx/build/html/python/_autosummary/mlx.core.distributed.init.html

Initializes the communication backend and creates the global communication group. It can be configured with strictness and a specific backend.

```APIDOC
## init()

### Description
Initialize the communication backend and create the global communication group.

### Signature
`init(_strict : bool = False_, _backend : str = 'any'_) → Group`

### Parameters
* **strict** (_bool_ _,__optional_) – If set to False it returns a singleton group in case `mx.distributed.is_available()` returns False otherwise it throws a runtime error. Default: `False`
* **backend** (_str_ _,__optional_) – Which distributed backend to initialize. Possible values `mpi`, `ring`, `nccl`, `jaccl`, `any`. If set to `any` all available backends are tried and the first one that succeeds becomes the global group which will be returned in subsequent calls. Default: `any`

### Returns
The group representing all the launched processes.

### Return type
_Group_

### Example
```python
import mlx.core as mx

group = mx.distributed.init(backend="ring")
```
```

--------------------------------

### device_count

Source: https://ml-explore.github.io/mlx/build/html/python/_autosummary/mlx.core.device_count.html

Get the number of available devices for the given device type.

```APIDOC
## device_count()

### Description
Get the number of available devices for the given device type.

### Parameters
#### Path Parameters
- **device_type** (mlx.core.DeviceType) - The type of device to query (cpu or gpu).

### Returns
Number of devices.

### Return Type
int
```

--------------------------------

### Initialize and Use Llama Model for Generation

Source: https://ml-explore.github.io/mlx/build/html/examples/llama-inference.html

Demonstrates initializing a Llama model, materializing its parameters, processing a prompt, and generating subsequent tokens. This snippet shows the end-to-end usage for inference.

```python
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)

```

--------------------------------

### mlx.core.arange

Source: https://ml-explore.github.io/mlx/build/html/python/_autosummary/mlx.core.arange.html

Generates numbers in the half-open interval [start, stop) in increments of step.

```APIDOC
## arange(start : Union[int, float], stop : Union[int, float], step : Union[None, int, float], dtype: Optional[Dtype] = None, *, stream: Union[None, Stream, Device] = None) -> array

### Description
Generates numbers in the half-open interval `[start, stop)` in increments of `step`.

### Parameters
#### Positional Parameters
- **start** (float or int, optional): Starting value which defaults to `0`.
- **stop** (float or int): Stopping value.
- **step** (float or int, optional): Increment which defaults to `1`.
#### Keyword Parameters
- **dtype** (Dtype, optional): Specifies the data type of the output. If unspecified will default to `float32` if any of `start`, `stop`, or `step` are `float`. Otherwise will default to `int32`.
- **stream** (Stream or Device, optional): Specifies the stream or device to use for the operation.

### Returns
- **array**: The range of values.

### Note
Following the Numpy convention the actual increment used to generate numbers is `dtype(start + step) - dtype(start)`. This can lead to unexpected results for example if start + step is a fractional value and the dtype is integral.
```

```APIDOC
## arange(stop : Union[int, float], step : Union[None, int, float] = None, dtype: Optional[Dtype] = None, *, stream: Union[None, Stream, Device] = None) -> array

### Description
Generates numbers in the half-open interval `[0, stop)` in increments of `step`.

### Parameters
#### Positional Parameters
- **stop** (float or int): Stopping value.
- **step** (float or int, optional): Increment which defaults to `1`.
#### Keyword Parameters
- **dtype** (Dtype, optional): Specifies the data type of the output. If unspecified will default to `float32` if any of `stop` or `step` are `float`. Otherwise will default to `int32`.
- **stream** (Stream or Device, optional): Specifies the stream or device to use for the operation.

### Returns
- **array**: The range of values.
```

--------------------------------

### slice

Source: https://ml-explore.github.io/mlx/build/html/cpp/ops.html

Creates a view of an array by slicing it using start, stop, and stride shapes.

```APIDOC
## slice

### Description
Create a view of an array by slicing it using start, stop, and stride shapes.

### Signature
```cpp
slice(const array &a, Shape start, Shape stop, Shape strides, StreamOrDevice s = {})
```
```

--------------------------------

### Create and Use Exponential Decay Scheduler

Source: https://ml-explore.github.io/mlx/build/html/python/optimizers/_autosummary/mlx.optimizers.exponential_decay.html

Demonstrates how to create an exponential decay learning rate scheduler and use it with an optimizer. The learning rate decreases over optimizer updates.

```python
>>> lr_schedule = optim.exponential_decay(1e-1, 0.9)
>>> optimizer = optim.SGD(learning_rate=lr_schedule)
>>> optimizer.learning_rate
array(0.1, dtype=float32)
>>>
>>> for _ in range(5): optimizer.update({}, {})
...
>>> optimizer.learning_rate
array(0.06561, dtype=float32)
```

--------------------------------

### MaxPool3d Layer Initialization and Usage

Source: https://ml-explore.github.io/mlx/build/html/python/nn/_autosummary/mlx.nn.MaxPool3d.html

Demonstrates how to initialize and use the MaxPool3d layer with a kernel size and stride. This is useful for reducing the spatial dimensions of 3D data.

```python
>>> import mlx.core as mx
>>> import mlx.nn.layers as nn
>>> x = mx.random.normal(shape=(8, 16, 32, 32, 4))
>>> pool = nn.MaxPool3d(kernel_size=2, stride=2)
>>> pool(x)
```

--------------------------------

### Instantiate and Use AvgPool1d

Source: https://ml-explore.github.io/mlx/build/html/python/nn/_autosummary/mlx.nn.AvgPool1d.html

Demonstrates how to instantiate the AvgPool1d layer with a specified kernel size and stride, and then apply it to an input tensor. Ensure mlx.core and mlx.nn.layers are imported.

```python
import mlx.core as mx
import mlx.nn.layers as nn
x = mx.random.normal(shape=(4, 16, 5))
pool = nn.AvgPool1d(kernel_size=2, stride=2)
pool(x)
```

--------------------------------

### slice

Source: https://ml-explore.github.io/mlx/build/html/python/ops.html

Extracts a sub-array from the input array based on start indices, axes, and slice size.

```APIDOC
## slice

### Description
Extract a sub-array from the input array.

### Signature
`slice(a, start_indices, axes, slice_size, *)`
```

--------------------------------

### Import and Use Custom MLX Extension

Source: https://ml-explore.github.io/mlx/build/html/dev/extensions.html

Demonstrates importing a custom MLX extension (axpby) and using it with MLX arrays. Shows how to perform element-wise operations with custom kernels.

```python
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 is correct: {mx.all(c == 6.0).item()}")

```

--------------------------------

### Create and Inspect MLX Arrays

Source: https://ml-explore.github.io/mlx/build/html/usage/quick_start.html

Import `mlx.core` and create arrays. Inspect their shapes and data types.

```python
import mlx.core as mx
a = mx.array([1, 2, 3, 4])
a.shape
a.dtype
b = mx.array([1.0, 2.0, 3.0, 4.0])
b.dtype
```

--------------------------------

### Check Python Processor

Source: https://ml-explore.github.io/mlx/build/html/install.html

Verify if your Python installation is native ARM. Run this command to check the processor type.

```python
python -c "import platform; print(platform.processor())"
```

--------------------------------

### Initialize and Use MaxPool2d

Source: https://ml-explore.github.io/mlx/build/html/python/nn/_autosummary/mlx.nn.MaxPool2d.html

Demonstrates initializing a MaxPool2d layer with a kernel size and stride, and then applying it to a random input tensor. Ensure mlx.core and mlx.nn.layers are imported.

```python
import mlx.core as mx
import mlx.nn.layers as nn
x = mx.random.normal(shape=(8, 32, 32, 4))
pool = nn.MaxPool2d(kernel_size=2, stride=2)
pool(x)
```

--------------------------------

### get_cache_memory()

Source: https://ml-explore.github.io/mlx/build/html/python/_autosummary/mlx.core.get_cache_memory.html

Get the cache size in bytes. The cache includes memory not currently used that has not been returned to the system allocator.

```APIDOC
## get_cache_memory()

### Description
Get the cache size in bytes. The cache includes memory not currently used that has not been returned to the system allocator.

### Method
`get_cache_memory()`

### Returns
- **int**: The cache size in bytes.
```

--------------------------------

### Configure JACCL Backend and Create Hostfile

Source: https://ml-explore.github.io/mlx/build/html/usage/launching_distributed.html

Use `mlx.distributed_config` to set up RDMA over Thunderbolt for JACCL. This command configures nodes, verifies connectivity, and generates a hostfile for distributed launches. Requires password-less sudo for auto-setup.

```bash
mlx.distributed_config --verbose --backend jaccl \
     --hosts m3-ultra-1,m3-ultra-2,m3-ultra-3,m3-ultra-4 --over thunderbolt \
     --auto-setup --output m3-ultra-jaccl.json
```

--------------------------------

### Create and Use a Linear Learning Rate Schedule

Source: https://ml-explore.github.io/mlx/build/html/python/optimizers/_autosummary/mlx.optimizers.linear_schedule.html

Demonstrates how to create a linear learning rate schedule and use it with an Adam optimizer. The learning rate increases linearly from 0 to 1e-1 over 100 steps, and remains at 1e-1 thereafter.

```python
>>> lr_schedule = optim.linear_schedule(0, 1e-1, 100)
>>> optimizer = optim.Adam(learning_rate=lr_schedule)
>>> optimizer.learning_rate
array(0.0, dtype=float32)
>>> for _ in range(101): optimizer.update({}, {})
... 
>>> optimizer.learning_rate
array(0.1, dtype=float32)
```

--------------------------------

### Install MLX in Debug Mode

Source: https://ml-explore.github.io/mlx/build/html/dev/metal_logging.html

Build MLX in debug mode to enable Metal logging. This is a prerequisite for kernel logging.

```bash
DEBUG=1 python -m pip install -e .
```

--------------------------------

### slice_update()

Source: https://ml-explore.github.io/mlx/build/html/python/_autosummary/mlx.core.slice_update.html

Updates a sub-array of the input array with the provided update array, starting at the specified indices along the given axes.

```APIDOC
## slice_update()

### Description
Updates a sub-array of the input array.

### Parameters
* **a** (_array_) – The input array to update.
* **update** (_array_) – The update array.
* **start_indices** (_array_) – The index location to start the slice at.
* **axes** (_tuple_ _(__int_ _)_) – The axes corresponding to the indices in `start_indices`.

### Returns
_array_ - The output array with the same shape and type as the input.

### Example
```python
>>> import mlx.core as mx
>>> a = mx.zeros((3, 3))
>>> mx.slice_update(a, mx.ones((1, 2)), start_indices=mx.array(1, 1), axes=(0, 1))
array([[0, 0, 0],
       [0, 1, 0],
       [0, 1, 0]], dtype=float32)
```
```

--------------------------------

### MultiOptimizer

Source: https://ml-explore.github.io/mlx/build/html/python/optimizers/_autosummary/mlx.optimizers.MultiOptimizer.html

Initializes the MultiOptimizer with a list of optimizers and their corresponding filters.

```APIDOC
## MultiOptimizer

### Description
Initializes the MultiOptimizer with a list of optimizers and their corresponding filters. The predicates take the full “path” of the weight and the weight itself and return True if it should be considered for this optimizer. The last optimizer in the list is a fallback optimizer and no predicate should be given for it.

### Parameters
* **optimizers** (_list_ _[__Optimizer_ _]_) – A list of optimizers to delegate to
* **filters** (_list_ _[__Callable_ _[__[__str_ _,__array_ _]__,__bool_ _]_) – A list of predicates that should be one less than the provided optimizers.
```

--------------------------------

### array slice(const array &a, Shape start, Shape stop, StreamOrDevice s = {})

Source: https://ml-explore.github.io/mlx/build/html/cpp/ops.html

Slice an array with a stride of 1 in each dimension.

```APIDOC
## array slice(const array &a, Shape start, Shape stop, StreamOrDevice s = {})

### Description
Slice an array with a stride of 1 in each dimension.

### Parameters
- **a** (const array &) - The input array.
- **start** (Shape) - The starting indices for slicing.
- **stop** (Shape) - The stopping indices for slicing.
- **s** (StreamOrDevice) - Optional stream or device for computation.
```

--------------------------------

### Evaluation Triggered by Control Flow

Source: https://ml-explore.github.io/mlx/build/html/usage/lazy_evaluation.html

Example where using a scalar array for control flow (if condition) implicitly triggers an evaluation.

```python
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

```

--------------------------------

### Lazy Model Initialization and Weight Loading

Source: https://ml-explore.github.io/mlx/build/html/usage/lazy_evaluation.html

Shows how model instantiation and weight loading are lazy. No memory is used until an explicit evaluation occurs.

```python
model = Model() # no memory used yet
model.load_weights("weights_fp16.safetensors")

```

--------------------------------

### Apply a function to all parameters

Source: https://ml-explore.github.io/mlx/build/html/python/nn/_autosummary/mlx.nn.Module.apply.html

Use Module.apply() to map a function across all parameters of a module. This example casts all parameters to 16-bit floats.

```python
model.apply(lambda x: x.astype(mx.float16))
```

--------------------------------

### MLX Standard Ops grid_sample Reference Implementation

Source: https://ml-explore.github.io/mlx/build/html/dev/custom_metal_kernels.html

Provides a reference implementation of grid_sample using standard MLX operations for comparison.

```python
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

```

--------------------------------

### init_single

Source: https://ml-explore.github.io/mlx/build/html/python/optimizers/_autosummary/mlx.optimizers.Adam.html

Initialize optimizer state for a single parameter.

```APIDOC
## init_single

### Description
Initialize optimizer state for a single parameter.

### Parameters
*   **parameter** - The parameter to initialize the state for.
*   **state** - The optimizer state dictionary.
```

--------------------------------

### ShardedToAllLinear.from_linear Classmethod

Source: https://ml-explore.github.io/mlx/build/html/python/nn/_autosummary/mlx.nn.ShardedToAllLinear.html

Converts a standard linear layer into a ShardedToAllLinear layer. This is useful for adapting existing models to a sharded computation setup.

```APIDOC
from_linear(linear_layer, *, segments, group)

Parameters:
    * **linear_layer** – The linear layer to convert.
    * **segments** – The number of segments to split the layer into.
    * **group** (_mx.distributed.Group_, optional) – The group to use for sharding. Defaults to the global group.
```

--------------------------------

### load

Source: https://ml-explore.github.io/mlx/build/html/python/ops.html

Loads array(s) from a binary file.

```APIDOC
## load

### Description
Load array(s) from a binary file.

### Signature
load(file, /[, format, return_metadata, stream])
```

--------------------------------

### Build Python API with CUDA Support

Source: https://ml-explore.github.io/mlx/build/html/install.html

Install the MLX Python API with CUDA support enabled by setting the CMAKE_ARGS environment variable.

```bash
CMAKE_ARGS="-DMLX_BUILD_CUDA=ON" pip install -e ".[dev]"

```

--------------------------------

### Build MLX Extension In-Place

Source: https://ml-explore.github.io/mlx/build/html/install.html

Build the MLX C++ extension in-place after installing development dependencies. This is useful for faster iteration during development.

```bash
python setup.py build_ext --inplace
```

--------------------------------

### Using Module with value_and_grad

Source: https://ml-explore.github.io/mlx/build/html/python/nn.html

Demonstrates how to use MLX modules with `mlx.core.value_and_grad`. Ensure parameters are passed as arguments to the function being transformed.

```python
model = ...

def f(params, other_inputs):
    model.update(params)  # <---- Necessary to make the model use the passed parameters
    return model(other_inputs)

f(model.trainable_parameters(), mx.zeros((10,)))
```

--------------------------------

### linspace()

Source: https://ml-explore.github.io/mlx/build/html/python/_autosummary/mlx.core.linspace.html

Generates `num` evenly spaced numbers over the interval `[start, stop]`. You can specify the number of samples, the data type, and the stream or device.

```APIDOC
## linspace()

### Description
Generate `num` evenly spaced numbers over interval `[start, stop]`.

### Parameters
* **start** (_scalar_) – Starting value.
* **stop** (_scalar_) – Stopping value.
* **num** (_int | None = 50_) – Number of samples, defaults to `50`.
* **dtype** (_Dtype | None = float32_) – Specifies the data type of the output, default to `float32`.
* **stream** (_None | Stream | Device = None_) – Specifies the stream or device for the operation.

### Returns
The range of values.

### Return type
array
```

--------------------------------

### Find Python for MLX CMake Integration

Source: https://ml-explore.github.io/mlx/build/html/dev/mlx_in_cpp.html

Configures CMake to find the Python interpreter and development components. This is necessary if MLX was installed via Python.

```cmake
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)

```

--------------------------------

### Compute Kronecker Product with mlx.core.kron

Source: https://ml-explore.github.io/mlx/build/html/python/_autosummary/mlx.core.kron.html

Demonstrates how to compute the Kronecker product of two arrays using mlx.core.kron. Ensure mlx is imported as mx.

```python
>>> a = mx.array([[1, 2], [3, 4]])
>>> b = mx.array([[0, 5], [6, 7]])
>>> result = mx.kron(a, b)
>>> print(result)
array([[0, 5, 0, 10],
       [6, 7, 12, 14],
       [0, 15, 0, 20],
       [18, 21, 24, 28]], dtype=int32)
```

--------------------------------

### Launch MPI Program with Explicit Library Paths

Source: https://ml-explore.github.io/mlx/build/html/usage/distributed.html

Manually launch an MPI program using `mpirun`, specifying the DYLD_LIBRARY_PATH and MPI_LIBNAME environment variables. This is an alternative to `mlx.launch`.

```bash
$ mpirun -np 2 -x DYLD_LIBRARY_PATH=/opt/homebrew/lib/ -x MPI_LIBNAME=libmpi.40.dylib python test.py
```