MLX

https://ml-explore.github.io/mlx/build/html/
Admin
MLX is a NumPy-like array framework designed for efficient and flexible machine learning on Apple...
Tokens:400,870
Snippets:3,347
Trust Score:8.3
Update:3 days ago
Show doc for...
Tokens:
Raw
### CMakeLists.txt basic setup

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

Basic CMake configuration for the example project. Save as CMakeLists.txt.

```cmake
cmake_minimum_required(VERSION 3.27)

project(example LANGUAGES CXX)

set(CMAKE_CXX_STANDARD 20)
set(CMAKE_CXX_STANDARD_REQUIRED ON)

```

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

### Run the example

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

Command to run the built example executable.

```bash
./build/example

```

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

### Build the example with CMake

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

Commands to configure and build the example project using CMake.

```bash
cmake -B build -DCMAKE_BUILD_TYPE=Release
cmake --build build

```

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

### Install MLX C++ Library

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

Install the compiled C++ library.

```bash
make install
```

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

### Install Linux Build Dependencies

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

Install required BLAS and LAPACK headers on Ubuntu.

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

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

### Install Xcode command line tools

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

Command to install the necessary Xcode developer tools.

```bash
xcode-select --install
```

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

### Setup parameters and load data

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

Configures hyperparameters and loads the MNIST dataset.

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

```

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

### Import MLX modules

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

Initial setup for using MLX core and neural network modules.

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

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

### Install MLX Python Library

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

Install the MLX Python library using pip.

```bash
pip install .
```

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

### Install MLX Python for Development

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

Install the package with development dependencies in editable mode.

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

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

### Backend Initialization Logic

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

Examples showing how different backend initialization calls affect subsequent calls to init().

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

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

### Build and Install MLX C++ Library

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

Create a build directory, configure with CMake, and compile the C++ library.

```bash
mkdir -p build && cd build
cmake .. && make -j
```

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

### Using glorot_uniform initializer

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

Demonstrates initializing an array with the glorot_uniform function, including an example with a custom gain parameter.

```python
>>> init_fn = nn.init.glorot_uniform()
>>> init_fn(mx.zeros((2, 2)))
array([[0.223404, -0.890597],
       [-0.379159, -0.776856]], dtype=float32)
>>> init_fn(mx.zeros((2, 2)), gain=4.0)
array([[-1.90041, 3.02264],
       [-0.912766, 4.12451]], dtype=float32)
```

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

### Exporting a function with mlx.export_function

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

This example demonstrates how to define a simple function and export it to a file named 'fun.mlxfn' using example array inputs.

```python
def fun(x, y):
    return x + y

x = mx.array(1)
y = mx.array([1, 2, 3])
mx.export_function("fun.mlxfn", fun, x, y=y)
```

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

### Install MLX with pip

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

Install the MLX Python package using pip. This is the first step before using MLX in a C++ project.

```bash
pip install -U mlx

```

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

### Install MLX on Apple silicon

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

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

```bash
pip install mlx
```

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

### Launch multi-node NCCL job

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

Example of launching a multi-node, multi-gpu job using the NCCL backend.

```bash
mlx.launch --backend nccl --hosts linux-1,linux-2 -n 8 -- ./my-job.sh
```

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

### Extracting sub-arrays with mx.slice

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

Examples demonstrating how to use mx.slice to extract sub-arrays from a 2D array by specifying start indices, axes, and slice sizes.

```python
>>> a = mx.array([[1, 2, 3], [4, 5, 6]])
>>> mx.slice(a, start_indices=mx.array(1), axes=(0,), slice_size=(1, 2))
array([[4, 5]], dtype=int32)
>>>
>>> mx.slice(a, start_indices=mx.array(1), axes=(1,), slice_size=(2, 1))
array([[2],
       [5]], dtype=int32)
```

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

### Updating a sub-array with slice_update

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

This example demonstrates how to initialize a 3x3 zero array and update a sub-section using slice_update with specified start indices and axes.

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

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

### Configure setup.py with CMakeExtension

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

Use the CMakeExtension and CMakeBuild classes from mlx.extension to define C++ and Metal extensions in a standard setuptools setup.py file.

```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")],
        cmdclass={"build_ext": extension.CMakeBuild},
        packages=[],
        extras_require={"dev":[]},
        zip_safe=False,
        python_requires=">=3.10",
    )
```

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

### Set MLX_ROOT for non-standard installation

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

Set MLX_ROOT to the installation location if MLX is installed to a non-standard path or CMake cannot find it.

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

```

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

### Initialize a sparse matrix with mlx.nn.init.sparse

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

Demonstrates how to create an initializer function with 50% sparsity and apply it to a zero-initialized array.

```python
>>> init_fn = nn.init.sparse(sparsity=0.5)
>>> init_fn(mx.zeros((2, 2)))
array([[-1.91187, 0],

> [0, -0.117483]], dtype=float32)
```

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

### slice (dynamic start)

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

Slices an array with dynamic starting indices.

```APIDOC
## slice (dynamic start)

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

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

### Parameters
- **a** (const array&) - Required - Input array.
- **start** (const array&) - Required - Dynamic starting indices.
- **axes** (std::vector<int>) - Required - Axes along which to slice.
- **slice_size** (Shape) - Required - Size of the slice.
- **s** (StreamOrDevice) - Optional - Stream or device for execution.

### Returns
- **array** - Sliced array.
```

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

### Import MLX and setup metadata

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

Initializes the MLX core package and defines hyperparameters for the linear regression model.

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

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

### Initialize and apply BatchNorm

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

Example showing how to import the module, initialize a BatchNorm layer with 4 features, and apply it to a random input tensor.

```python
>>> import mlx.core as mx
>>> import mlx.nn as nn
>>> x = mx.random.normal((5, 4))
>>> bn = nn.BatchNorm(num_features=4, affine=True)
>>> output = bn(x)
```

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

### slice_update (dynamic start)

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

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

```APIDOC
## slice_update (dynamic start)

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

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

### Parameters
- **src** (const array&) - Required - Source array.
- **update** (const array&) - Required - Update array.
- **start** (const array&) - Required - Dynamic starting indices.
- **axes** (std::vector<int>) - Required - Axes along which to update.
- **s** (StreamOrDevice) - Optional - Stream or device for execution.

### Returns
- **array** - Updated array.
```

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

### Install OpenMPI via Conda

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

Install the OpenMPI library using the Anaconda package manager.

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

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

### Run training loop

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

Instantiates the model and optimizer, then executes the training loop.

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

```

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

### Manual MPI configuration with mpirun

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

Manually configure MPI library paths using environment variables, or use mlx.launch to automate the process.

```bash
$ mpirun -np 2 -x DYLD_LIBRARY_PATH=/opt/homebrew/lib/ -x MPI_LIBNAME=libmpi.40.dylib python test.py
$ # or simply
$ mlx.launch -n 2 test.py
```

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

### Install CUDA Build Dependencies

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

Install CUDA toolkit and required headers on Ubuntu for 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
```

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

### Install MLX for CPU-only on Linux

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

Install the CPU-only version of MLX for Linux. Requires glibc >= 2.35 and Python >= 3.10.

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

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

### Initialize an array with uniform distribution

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

This example demonstrates creating an initializer function with specified bounds and applying it to a zero-filled array.

```python
>>> init_fn = nn.init.uniform(low=0, high=1)
>>> init_fn(mx.zeros((2, 2)))
array([[0.883935, 0.863726],
       [0.617261, 0.417497]], dtype=float32)
```

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

### linspace(start, stop[, num, endpoint, ...])

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

Generates num evenly spaced numbers over interval [start, stop].

```APIDOC
## linspace(start, stop[, num, endpoint, ...])

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

### Parameters
- **start** (scalar) - Start value.
- **stop** (scalar) - Stop value.
- **num** (int, optional) - Number of samples (default 50).
- **endpoint** (bool, optional) - Whether to include stop (default True).
- **...** - Additional parameters (not specified).

### Returns
- **array** - Evenly spaced numbers.
```

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

### Initialize an array with nn.init.uniform

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

Demonstrates creating an initializer function and applying it to a specific array.

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

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

### Perform QR factorization with mx.linalg.qr

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

Example showing the QR factorization of a 2x2 matrix using the CPU stream.

```python
>>> A = mx.array([[2., 3.], [1., 2.]])
>>> Q, R = mx.linalg.qr(A, stream=mx.cpu)
>>> Q
array([[-0.894427, -0.447214],
       [-0.447214, 0.894427]], dtype=float32)
>>> R
array([[-2.23607, -3.57771],
       [0, 0.447214]], dtype=float32)
```

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

### init

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

Initialize the communication backend and create the global communication group.

```APIDOC
## init

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

### Parameters
- **strict** (optional) - Whether to enforce strict initialization.
- **backend** (optional) - The backend to initialize.
- **all_gather_factory** (optional) - Factory for all-gather operations.
```

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

### Find MLX when installed with Python

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

Add this to the CMake file if MLX was installed with Python. It locates the MLX CMake directory using the Python interpreter.

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

```

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

### Initialize weights with glorot_normal

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

Demonstrates creating a Glorot normal initializer and applying it to arrays with different gain values.

```python
>>> init_fn = nn.init.glorot_normal()
>>> init_fn(mx.zeros((2, 2)))
array([[0.191107, 1.61278],
       [-0.150594, -0.363207]], dtype=float32)
>>> init_fn(mx.zeros((2, 2)), gain=4.0)
array([[1.89613, -4.53947],
       [4.48095, 0.995016]], dtype=float32)
```

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

### Cross Entropy Loss Examples

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

Examples demonstrating the use of cross_entropy with class indices, probability targets, and handling of different data types like bfloat16.

```python
>>> import mlx.core as mx
>>> import mlx.nn as nn
>>>
>>> # Class indices as targets
>>> logits = mx.array([[2.0, -1.0], [-1.0, 2.0]])
>>> targets = mx.array([0, 1])
>>> nn.losses.cross_entropy(logits, targets)
array([0.0485873, 0.0485873], dtype=float32)
>>>
>>> # Probabilities (or one-hot vectors) as targets
>>> logits = mx.array([[2.0, -1.0], [-1.0, 2.0]])
>>> targets = mx.array([[0.9, 0.1], [0.1, 0.9]])
>>> nn.losses.cross_entropy(logits, targets)
array([0.348587, 0.348587], dtype=float32)
>>>
>>> # Half precision logits with class indices as targets. On the GPU a
>>> # fused kernel accumulates the reduction in float32:
>>> logits = mx.array([[2.0, -1.0], [-1.0, 2.0]], mx.bfloat16)
>>> targets = mx.array([0, 1])
>>> nn.losses.cross_entropy(logits, targets)
array([0.048584, 0.048584], dtype=bfloat16)
>>>
>>> # The CPU reduces in the dtype of the logits, so upcast them to get
>>> # the same accuracy:
>>> nn.losses.cross_entropy(logits.astype(mx.float32), targets)
array([0.0485873, 0.0485873], dtype=float32)
```

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

### Lion.__init__

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

Constructor for the Lion optimizer. Initializes the optimizer with learning rate, betas, and weight decay.

```APIDOC
## Lion.__init__

### Description
Initializes the Lion optimizer with the given learning rate, betas, and weight decay.

### Method
Constructor

### Signature
Lion(learning_rate: float | Callable[[array], array], betas: List[float] = [0.9, 0.99], weight_decay: float = 0.0)

### Parameters
- **learning_rate** (float or callable) - Required - The learning rate eta.
- **betas** (Tuple[float, float], optional) - The coefficients (beta1, beta2) used for computing the gradient momentum and update direction. Default: (0.9, 0.99)
- **weight_decay** (float, optional) - The weight decay lambda. Default: 0.0

### Notes
The Lion optimizer uses sign operations for updates, so a learning rate 3-10x smaller than AdamW and a weight decay 3-10x larger than AdamW is recommended to maintain the strength (lr * wd).
```

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

### Initialize Optimizer State

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

Demonstrates initializing an SGD optimizer with model parameters and accessing the resulting state keys.

```python
>>> optimizer = optim.SGD(learning_rate=1e-1, momentum=0.9)
>>> model = nn.Linear(2, 2)
>>> optimizer.init(model.trainable_parameters())
>>> optimizer.state.keys()
dict_keys(['step', 'learning_rate', 'weight', 'bias'])
```

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

### arange(start, stop, step, dtype, *, stream)

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. This is the first overload of the arange function.

```APIDOC
## arange(start, stop, step, dtype, *, stream)

### Description
Generates numbers in the half-open interval [start, stop) in increments of step. This is the first overload of the arange function.

### Signature
arange(start: int | float, stop: None | int | float, step: None | int | float, dtype: Dtype | None = None, *, stream: StreamOrDevice = None) -> array

### Parameters
- **start** (int | float) - Required - Starting value which defaults to 0.
- **stop** (int | float) - Required - Stopping value.
- **step** (int | float) - Required - Increment which defaults to 1.
- **dtype** (Dtype, optional) - 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, or int64 if any of start, stop, or step does not fit in int32.
- **stream** (StreamOrDevice, keyword-only) - Optional - The stream or device on which to execute the operation.

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

### Notes
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.
```

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

### Install MLX with CUDA 13 support

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

Install the MLX package with CUDA 13 backend. Requires Nvidia driver >= 580 or an appropriate CUDA compatibility package.

```bash
pip install mlx[cuda13]
```

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

### Using step_decay with an optimizer

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

Demonstrates initializing a step decay schedule and applying it to an SGD optimizer, showing the learning rate update after 21 steps.

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

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

### Example: Using he_uniform initializer

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

Shows how to create a He uniform initializer and apply it to a zero array. The first call uses the default fan_in mode and gain of 1, while the second uses mode='fan_out' and gain=5.

```python
>>> init_fn = nn.init.he_uniform()
>>> init_fn(mx.zeros((2, 2)))  # uses fan_in
array([[0.0300242, -0.0184009],
       [0.793615, 0.666329]], dtype=float32)
>>> init_fn(mx.zeros((2, 2)), mode="fan_out", gain=5)
array([[-1.64331, -2.16506],
       [1.08619, 5.79854]], dtype=float32)
```

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

### Mixed CPU/GPU computation example

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

A contrived example showing how to assign different parts of a computation to different devices. The matmul runs on the GPU (d1) and the repeated exp operations run on the CPU (d2).

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

```python
a = mx.random.uniform(shape=(4096, 512))
b = mx.random.uniform(shape=(512, 4))
```

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

### device_info([d])

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

Get information about a device.

```APIDOC
## device_info([d])

### Description
Get information about a device.

### Method
Function

### Parameters
- **d** (optional) - The device to get information about. If not provided, uses the default device.

### Returns
Information about the device.
```

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

### MaxPool2d usage example

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

Shows how to create a MaxPool2d layer with kernel_size=2 and stride=2, and apply it to a random input tensor of shape (8, 32, 32, 4). The example imports mlx.core as mx and mlx.nn.layers as nn.

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

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

### Device.__init__(self, type: mlx.core.DeviceType, index: int = 0) -> None

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

Initializes a new Device instance with the specified device type and optional index. This constructor is used to create a device object for running operations.

```APIDOC
## Device.__init__(self, type: mlx.core.DeviceType, index: int = 0) -> None

### Description
Initializes a new Device instance with the specified device type and optional index.

### Method
Constructor

### Signature
`__init__(self, type: mlx.core.DeviceType, index: int = 0) -> None`

### Parameters
- **type** (mlx.core.DeviceType) - Required - The type of the device.
- **index** (int) - Optional - The device index. Defaults to 0.

### Returns
None
```

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

### Install MLX with CUDA 12 support

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

Install the MLX package with CUDA 12 backend. Requires Nvidia architecture >= SM 7.5, driver >= 550.54.14, CUDA toolkit >= 12.0, glibc >= 2.35, and Python >= 3.10.

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

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

### Optimizer.init(parameters)

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

Initializes the optimizer's state for the given parameters. This method should be called before applying gradients.

```APIDOC
## Optimizer.init(parameters)

### Description
Initializes the optimizer's state for the given parameters.

### Method
init

### Parameters
- **parameters** - The parameters for which to initialize the optimizer state.

### Returns
- None (state is stored internally).
```

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

### Using scaled_dot_product_attention

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

Example demonstrating the initialization of query, key, and value tensors and the invocation of the scaled_dot_product_attention function with a causal mask.

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

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

### mlx.core.linspace(start, stop, num=50, endpoint=True, dtype=float32, stream=None) -> array

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]. The endpoint parameter controls whether stop is included, and dtype specifies the output data type.

```APIDOC
## mlx.core.linspace

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

### Function Signature
`linspace(start: scalar, stop: scalar, num: int | None = 50, endpoint: bool = True, dtype: Dtype | None = float32, stream: StreamOrDevice = None) -> array`

### Parameters
- **start** (scalar) - Required - Starting value.
- **stop** (scalar) - Required - Stopping value.
- **num** (int, optional) - Number of samples, defaults to `50`.
- **endpoint** (bool, optional) - If `True`, `stop` is the last sample. Otherwise it is not included and the samples are spaced over the half-open interval `[start, stop)`. Default: `True`.
- **dtype** (Dtype, optional) - Specifies the data type of the output, default to `float32`.
- **stream** (StreamOrDevice, optional) - Stream or device to use.

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

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

### iinfo(*args, **kwargs)

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

Get information on integer types.

```APIDOC
## iinfo(*args, **kwargs)

### Description
Get information on integer types.

### Parameters
- **args** (variadic) - Required - Arguments specifying the integer type.
- **kwargs** (variadic) - Optional - Additional keyword arguments.

### Returns
- **object** - Information about the integer type.
```

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

### Add executable and link MLX

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

Add the example.cpp program as an executable and link it against the mlx library.

```cmake
add_executable(example example.cpp)
target_link_libraries(example PRIVATE mlx)

```

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

### Usage of linear_schedule with an optimizer

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

Demonstrates initializing a linear schedule and applying it to an Adam optimizer, showing the learning rate progression over 100 steps.

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

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

### finfo(*args, **kwargs)

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

Get information on floating-point types.

```APIDOC
## finfo(*args, **kwargs)

### Description
Get information on floating-point types.

### Parameters
- **args** (variadic) - Required - Arguments specifying the floating-point type.
- **kwargs** (variadic) - Optional - Additional keyword arguments.

### Returns
- **object** - Information about the floating-point type.
```

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

### slice (strides)

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

Slices an array using start, stop, and strides.

```APIDOC
## slice (strides)

### Description
Slices an array using start, stop, and strides.

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

### Parameters
- **a** (const array&) - Required - Input array.
- **start** (Shape) - Required - Starting indices.
- **stop** (Shape) - Required - Stopping indices.
- **strides** (Shape) - Required - Strides for each dimension.
- **s** (StreamOrDevice) - Optional - Stream or device for execution.

### Returns
- **array** - Sliced array.
```

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

### Scatter index broadcasting

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

Example showing how index arrays are broadcasted together.

```cpp
auto indices = std::vector{array({2, 0}), array({1})};
```

```cpp
auto indices = std::vector{array({2, 0}), array({1, 1})};
```
Tokens:
Raw