MLX

Warnings

• gotcha MLX is primarily optimized for Apple Silicon. While CUDA backend support is in active development, performance and feature parity on non-Apple hardware (e.g., NVIDIA GPUs on Linux) might be limited or in beta. Expect potential performance differences and missing operations.
  Fix: Ensure you are running on Apple Silicon for best performance, or monitor MLX's GitHub releases for updates on CUDA backend stability and feature completeness if targeting NVIDIA GPUs. Report issues if you encounter missing operations or performance cliffs.
• gotcha `mlx.nn.Conv2d` expects input images in NHWC (batch, height, width, channels) format, which differs from the NCHW (batch, channels, height, width) format commonly used by frameworks like PyTorch. This requires transposing input tensors when porting models or using data loaders that provide NCHW data.
  Fix: Reshape your input tensors to NHWC format before passing them to `mlx.nn.Conv2d`. For example, a PyTorch-style NCHW tensor `x` can be converted using `x.transpose(0, 2, 3, 1)`.
• gotcha MLX uses lazy computation, meaning operations build a computation graph but do not execute immediately. Arrays are only materialized when their values are needed (e.g., printed, converted to NumPy, or explicitly evaluated with `mx.eval()`). Users expecting immediate execution might be surprised by this behavior.
  Fix: To force immediate computation and materialize an array, use `mx.eval(array)`. Be mindful of when and where you introduce `mx.eval()` to balance performance and debugging needs.
• gotcha When using function transformations like `mlx.core.value_and_grad()`, they operate on 'pure' functions. This means that model parameters or other state should be explicitly passed as arguments to the function being transformed, rather than relying on global mutable state or attributes, to ensure correct gradient computation and graph optimization.
  Fix: Refactor functions intended for `mlx.core` transformations to accept all necessary parameters as explicit inputs. For neural networks, `mlx.nn.value_and_grad()` can simplify this by handling trainable parameters of a `Module`.
• gotcha For optimal performance on Apple Silicon, it is crucial to use a native `arm` Python environment. Running MLX in an `x86` Python environment via Rosetta emulation will result in significantly degraded performance.
  Fix: Verify your Python environment using `uname -p`. If it returns `i386` or `x86_64` on an Apple Silicon Mac, switch to a native `arm` Python installation, typically managed via tools like `conda` or directly installing an `arm64` Python version.

Install

• pip install mlx Install MLX

Imports

• mlx.core

  import mlx.core as mx

  Provides core array operations, similar to NumPy.
• mlx.nn

  import mlx.nn as nn

  Offers neural network components like layers and activation functions, akin to PyTorch's `torch.nn`.
• mlx.optimizers

  import mlx.optimizers as optim

  Contains optimization algorithms for training neural networks.

Quickstart

This quickstart demonstrates creating MLX arrays, performing basic element-wise and matrix operations, and explicitly materializing computations using `mx.eval()` due to MLX's lazy evaluation.

import mlx.core as mx

# Create an MLX array
a = mx.array([1, 2, 3], mx.float32)
print(f"Array a: {a}, dtype: {a.dtype}, shape: {a.shape}")

# Perform an operation (e.g., exponential)
b = mx.exp(a)
print(f"Array b (exp(a)): {b}")

# Perform a matrix multiplication
c = mx.array([[1, 2], [3, 4]])
d = mx.array([[5, 6], [7, 8]])
e = mx.matmul(c, d)
print(f"Matrix c: {c}\nMatrix d: {d}\nMatrix product c @ d: {e}")

# Ensure computation is materialized (for lazy operations)
mx.eval(e)
print("Computation materialized.")