Learning MLX

I’ve only recenly came across MLX, Apple’s array framework. In comparison to JAX, my usual framework of choice, MLX features a proper C++ API making it particularly interesting. To explore the library a bit, I’ve implemented one of our recent publications, SABC (Albert et al., 2025), in MLX using its Python and C++ API.

You can find the SABC code on GitHub.

Why MLX

MLX is a C++ library that features a fairly high-level API and which exposes Python bindings. JAX is Python first, and anything C++ has to be done via Pallas or a custom FFI call. That is, in MLX we can write compute-itensive code, like an MCMC sampler, in C++ and expose it to Python through nanobind.

Conveniently when one has a JAX background, most of the Python API maps one to one onto JAX:

JAX’s control-flow primitives (lax.fori_loop, lax.scan, …) are necessary because a jitted function compiles to one static graph. A Python for or if on a traced value cannot go inside that graph, so the loop or the branch has to be written as a primitive the tracer can capture.

For instance, let’s look at Newton’s method in JAX:

from jax import lax
import jax.numpy as jnp

def newton(f, df, x, n_iter=20):
  def body(_, x):
    step = jnp.where(jnp.abs(df(x)) > 1e-12, f(x) / df(x), 0.0)
    return x - step
  return lax.fori_loop(0, n_iter, body, x)

MLX has no trace, so the same function can be done with an ordinary for loop using MLX’s Python bindings:

import mlx.core as mx

def newton(f, df, x, n_iter=20):
  for _ in range(n_iter):
    step = mx.where(mx.abs(df(x)) > 1e-12, f(x) / df(x), 0.0)
    x = x - step
  return x

The same loop can be transcribed straight into C++:

mx::array newton(const std::function<mx::array(mx::array)>& f,
                 const std::function<mx::array(mx::array)>& df,
                 mx::array x, int n_iter = 20) {
  for (int i = 0; i < n_iter; ++i) {
    mx::array step = mx::where(
        mx::greater(mx::abs(df(x)), mx::array(1e-12f)),
        mx::divide(f(x), df(x)), mx::array(0.0f));
    x = mx::subtract(x, step);
  }
  return x;
}

The other primitives translate the same way. A lax.scan is a loop that accumulates. For instance, let’s consider exponential moving average in JAX:

def ema(xs, alpha):
  def step(carry, x):
    y = alpha * x + (1 - alpha) * carry
    return y, y
  _, ys = lax.scan(step, xs[0], xs)
  return ys

In MLX the carry is a plain variable and the scan is a plain loop:

def ema(xs, alpha):
  carry = xs[0]
  ys = []
  for x in xs:
    carry = alpha * x + (1 - alpha) * carry
    ys.append(carry)
  return mx.stack(ys)

Each MLX op is lazy, so a loop builds a graph that keeps growing until mx.eval evaluates it (i.e., actually computes a value). However, similarly to JAX, mx.compile can still compile a function into a single graph when high throughput is needed.

The benchmark

I’ve compared SABC in MLX against a JAX version, and a NumPy version and a Numba version from a collaborator: sbijax, sabc-mlx, sabc-numpy, sabc-numba. The table below shows wall time, compile time, peak RSS, and W₁ distance to an MCMC reference posterior on the two_moons task (the best value in each column is in bold).

All four implementations recover the reference posterior faithuflly. Post-compile, JAX runs about 4x faster than MLX. MLX wins on memory: 10x less, because there is no XLA buffer pool and no JIT cache.

Conclusion

In summary, even when factoring in the compile time I think I am still sticking to JAX. But I will definitely use MLX more often in the future 👾🍏.

References

Albert, C., Ulzega, S., Dirmeier, S., Scheidegger, A., Bassi, A., and Mira, A. (2025). Simulated Annealing ABC with multiple summary statistics. arXiv preprint arXiv:2505.23261.