mlx-rs-burn 0.25.5

Unofficial rust wrapper for Apple's mlx machine learning library (fork with slice, scatter, flip operations for burn-mlx)
docs.rs failed to build mlx-rs-burn-0.25.5
Please check the build logs for more information.
See Builds for ideas on how to fix a failed build, or Metadata for how to configure docs.rs builds.
If you believe this is docs.rs' fault, open an issue.

mlx-rs-burn

Crates.io Documentation License

Rust bindings for Apple's MLX framework — a fork of mlx-rs with additional operations for burn-mlx.

Why This Fork?

This crate extends the original mlx-rs with operations required by the burn-mlx deep learning backend:

  • slice - Array slicing with start/stop indices
  • slice_update - In-place slice updates
  • scatter / scatter_add - Scatter operations for gradient computation
  • flip - Array flipping along axes

These operations are essential for implementing neural network backward passes (e.g., pooling gradients).

Installation

[dependencies]
mlx-rs-burn = "0.25.4"

Quick Start

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]

Features

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

Feature Flags

  • metal (default) - Enable Metal GPU support
  • accelerate (default) - Enable Accelerate framework
  • safetensors - Enable safetensors file format support

Requirements

  • macOS with Apple Silicon (M1/M2/M3/M4)
  • Rust 1.82+

Related Crates

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

Upstream

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.

License

MIT OR Apache-2.0

Acknowledgments

  • mlx-rs - Original Rust bindings by OxideAI
  • MLX - Apple's machine learning framework
  • Burn - Rust deep learning framework