Rust bindings for Apple's MLX framework — a fork of mlx-rs with additional operations for burn-mlx.
This crate extends the original mlx-rs with operations required by the burn-mlx deep learning backend:
slice- Array slicing with start/stop indicesslice_update- In-place slice updatesscatter/scatter_add- Scatter operations for gradient computationflip- Array flipping along axes
These operations are essential for implementing neural network backward passes (e.g., pooling gradients).
[dependencies]
mlx-rs-burn = "0.25.4"use mlx_rs::Array;
use mlx_rs::ops::{slice, scatter_add, flip};
// Create an array
let x = Array::from_slice(&[1.0f32, 2.0, 3.0, 4.0], &[4]);
// Slice operations
let sliced = slice(&x, &[1], &[3], None).unwrap();
// Result: [2.0, 3.0]
// Scatter add (useful for gradient accumulation)
let zeros = Array::zeros::<f32>(&[4]);
let indices = Array::from_slice(&[0i32, 2], &[2]);
let updates = Array::from_slice(&[10.0f32, 30.0], &[2]);
let result = scatter_add(&zeros, &[&indices], &updates, &[0]).unwrap();
// Result: [10.0, 0.0, 30.0, 0.0]
// Flip along axis
let flipped = flip(&x, &[0]).unwrap();
// Result: [4.0, 3.0, 2.0, 1.0]All features from the original mlx-rs are available:
- Performance: Optimized for Apple Silicon (M1/M2/M3/M4)
- Lazy Evaluation: Arrays are only materialized when needed
- Dynamic Graphs: Computation graphs constructed dynamically
- Unified Memory: Zero-copy data sharing between CPU and GPU
- Metal GPU: Native GPU acceleration via Metal
metal(default) - Enable Metal GPU supportaccelerate(default) - Enable Accelerate frameworksafetensors- Enable safetensors file format support
- macOS with Apple Silicon (M1/M2/M3/M4)
- Rust 1.82+
| Crate | Description |
|---|---|
| burn-mlx | MLX backend for the Burn deep learning framework |
| mlx-sys-burn | Low-level FFI bindings |
| mlx-macros-burn | Procedural macros |
| mlx-internal-macros-burn | Internal macros |
This is a fork of oxideai/mlx-rs. We aim to contribute these additions back upstream. In the meantime, this fork is published to enable burn-mlx on crates.io.
MIT OR Apache-2.0