mlx_c_rust/
lib.rs

1//! Safe ownership boundary around the official MLX C API.
2//!
3//! This crate owns exactly one responsibility: the MLX-C C ABI translated
4//! into idiomatic Rust. It holds the bindgen-generated raw declarations, the
5//! captured-error machinery, the owned handle types (`MlxArray`, `MlxDtype`,
6//! `MlxArrayVector`, `MlxStream`), the per-worker bindings context
7//! (`MlxBindingsContext`: GPU stream and linked version), and every operation
8//! wrapper family: creation, shape, padding, elementwise math, activation,
9//! normalization, convolution, random, rope, quantized operations,
10//! quantization construction, NVFP4, attention, the compiled-graph family,
11//! the Metal kernels and capture, the general operation set, and the gap
12//! shim family — this repository's own `mlxrs_*` C ABI (compiled from
13//! `native_gap_shim/` by the native build engine) for MLX functions the
14//! vendored MLX-C pin has not wrapped yet, so an MLX capability is
15//! reachable the day the pin moves. It contains
16//! no embedding policy: memory limits, the metallib path selection and
17//! verification, safetensors vtable policies, and process I/O accounting
18//! belong to the embedding application, which links this crate against an
19//! MLX-C library and converts `MlxCError` values into its own typed runtime
20//! errors at its boundary.
21//!
22//! # API-surface completeness
23//!
24//! The bridge is complete, not selective: every public function and type the
25//! vendored MLX-C headers declare is carried here. The build script
26//! translates the vendored tree in full and emits a bridged-symbol inventory
27//! alongside the declarations, so any upstream pin bump that adds, removes,
28//! or renames surface shows up as an inventory change the coverage contract
29//! can fail on. The recorded state for the current pin lives in this crate's
30//! `COVERAGE_INVENTORY.md`.
31//!
32//! Source-reference audits must never drive removals here: a wrapper with no
33//! caller yet is the expected steady state of a complete bridge — the crate
34//! exists precisely so that every MLX-C capability is reachable without
35//! waiting for a wrapper to be written on demand.
36//!
37//! # Bindings-only mode
38//!
39//! Consumers who provide their own native MLX build — for example a pinned,
40//! patched fork produced by their own native orchestration — can take this
41//! crate as the bridge alone: depend with `default-features = false` and
42//! the build script emits no link directives at all. The vendored-header
43//! declarations, the captured-error machinery, the owned handle types, and
44//! every operation wrapper family are identical; the consumer's own build
45//! supplies and links the native libraries. The default `link` feature
46//! carries the whole discovery ladder for everyone else.
47//!
48//! # Test placement
49//!
50//! Tests that never touch MLX-C symbols (the error-channel machinery) live
51//! in this crate's test tree. Anything exercising the bindings themselves
52//! needs an MLX-C library at link time, discovered through the
53//! `MLX_C_RUST_MLX_C_PREFIX` environment variable or `pkg-config mlx-c`
54//! (see the build script and README); without one, library compilation and
55//! documentation builds still succeed and only test binaries fail to link.
56
57mod bindings;
58mod compile;
59mod metal;
60mod operations;
61mod runtime;
62mod types;
63
64pub mod error;
65pub mod raw;
66
67pub use bindings::context::MlxBindingsContext;
68pub use compile::elementwise_graphs::MlxCompiledElementwiseGraphs;
69pub use compile::gated_delta_rate_activation::MlxCompiledGatedDeltaRateActivation;
70pub use compile::graph::{
71    MlxGraphBuilder, array_from_vector, graph_output_array, set_graph_output,
72    set_graph_output_vector,
73};
74pub use compile::multi_output_graph::MlxCompiledMultiOutputGraph;
75pub use compile::norm_fusion::MlxCompiledDualHeadNorm;
76pub use compile::swiglu::MlxCompiledSwiGlu;
77pub use compile::transforms::{
78    MlxCompileCache, disable_compile, enable_compile, jvp_of_closure, set_compile_mode,
79    value_and_grad_of_closure, vjp_of_closure, vmap_replace, vmap_trace,
80};
81pub use error::{
82    MlxCError, check_status, clear_captured_mlx_error, install_non_terminating_error_handler,
83    take_captured_mlx_error,
84};
85pub use metal::cuda_kernel::{MlxCudaKernel, MlxCudaKernelConfig};
86pub use metal::kernel::{
87    MlxMetalKernel, MlxMetalKernelConfig, MlxMetalKernelOutput, MlxMetalKernelTemplateArgument,
88    apply_metal_kernel_in_graph,
89};
90pub use metal::metallib::set_metallib_path;
91pub use runtime::export::{
92    GraphOutputFile, MlxFunctionExporter, MlxImportedFunction, MlxNodeNamer, export_function,
93    export_function_with_keywords, export_graph_to_dot, print_graph,
94};
95pub use runtime::io::{
96    MlxIoGguf, MlxIoReader, MlxIoWriter, load_array_from_reader, load_safetensors,
97    load_safetensors_from_reader, save_array_to_writer, save_safetensors,
98    save_safetensors_to_writer,
99};
100pub use runtime::memory::{
101    active_memory_bytes, cache_memory_bytes, clear_allocator_cache, memory_limit_bytes,
102    peak_memory_bytes, reset_peak_memory, set_cache_limit, set_memory_limit, synchronize,
103};
104pub use runtime::platform::{
105    MlxDistributedGroup, cuda_is_available, distributed_is_available, initialize_distributed,
106    metal_is_available, metallib_path, set_wired_limit,
107};
108pub use runtime::stream_defaults::{
109    MlxStreamThreadLocal, clear_streams, set_default_stream, synchronize_all_defaults,
110};
111pub use types::array::{MlxArray, MlxDtype};
112pub use types::array_vector::MlxArrayVector;
113pub use types::closure_custom::{MlxClosureCustom, MlxClosureCustomJvp, MlxClosureCustomVmap};
114pub use types::closures::{MlxClosure, MlxClosureKwargs, MlxClosureValueAndGrad};
115pub use types::device::{MlxDeviceHandle, MlxDeviceInfo};
116pub use types::maps::{
117    MlxMapStringToArray, MlxMapStringToArrayIterator, MlxMapStringToString,
118    MlxMapStringToStringIterator,
119};
120pub use types::stream::MlxStream;
121pub use types::string::MlxString;
122pub use types::vectors::{MlxVectorInt, MlxVectorStream, MlxVectorString, MlxVectorVectorArray};
mlx_c_rust/lib.rs

mlx_c_rust/
lib.rs
