Skip to main content

Crate mlx_c_rust

Crate mlx_c_rust 

Source
Expand description

Safe ownership boundary around the official MLX C API.

This crate owns exactly one responsibility: the MLX-C C ABI translated into idiomatic Rust. It holds the bindgen-generated raw declarations, the captured-error machinery, the owned handle types (MlxArray, MlxDtype, MlxArrayVector, MlxStream), the per-worker bindings context (MlxBindingsContext: GPU stream and linked version), and every operation wrapper family: creation, shape, padding, elementwise math, activation, normalization, convolution, random, rope, quantized operations, quantization construction, NVFP4, attention, the compiled-graph family, the Metal kernels and capture, and the general operation set. It contains no embedding policy: memory limits, the metallib path selection and verification, safetensors vtable policies, and process I/O accounting belong to the embedding application, which links this crate against an MLX-C library and converts MlxCError values into its own typed runtime errors at its boundary.

§API-surface completeness

The bridge is complete, not selective: every public function and type the vendored MLX-C headers declare is carried here. The build script translates the vendored tree in full and emits a bridged-symbol inventory alongside the declarations, so any upstream pin bump that adds, removes, or renames surface shows up as an inventory change the coverage contract can fail on. The recorded state for the current pin lives in this crate’s COVERAGE_INVENTORY.md.

Source-reference audits must never drive removals here: a wrapper with no caller yet is the expected steady state of a complete bridge — the crate exists precisely so that every MLX-C capability is reachable without waiting for a wrapper to be written on demand.

§Bindings-only mode

Consumers who provide their own native MLX build — for example a pinned, patched fork produced by their own native orchestration — can take this crate as the bridge alone: depend with default-features = false and the build script emits no link directives at all. The vendored-header declarations, the captured-error machinery, the owned handle types, and every operation wrapper family are identical; the consumer’s own build supplies and links the native libraries. The default link feature carries the whole discovery ladder for everyone else.

§Test placement

Tests that never touch MLX-C symbols (the error-channel machinery) live in this crate’s test tree. Anything exercising the bindings themselves needs an MLX-C library at link time, discovered through the MLX_C_RUST_MLX_C_PREFIX environment variable or pkg-config mlx-c (see the build script and README); without one, library compilation and documentation builds still succeed and only test binaries fail to link.

Re-exports§

pub use error::MlxCError;
pub use error::check_status;
pub use error::clear_captured_mlx_error;
pub use error::install_non_terminating_error_handler;
pub use error::take_captured_mlx_error;

Modules§

error
Captured MLX-C error translation.
raw
Bindgen-generated declarations for the complete official MLX C boundary.

Structs§

GraphOutputFile
A rendered-graph destination backed by a C FILE* the caller owns for the duration of one rendering call.
MlxArray
Owned MLX array handle released exactly once through the official C API.
MlxArrayVector
Temporary MLX vector that retains array handles for one aggregate operation.
MlxBindingsContext
Process-global official MLX C bindings context configured for one worker.
MlxClosure
Owned MLX closure handle released exactly once through the C API.
MlxClosureCustom
Owned MLX custom-transform (VJP) closure, released exactly once.
MlxClosureCustomJvp
Owned MLX custom JVP closure, released exactly once.
MlxClosureCustomVmap
Owned MLX custom VMAP closure, released exactly once.
MlxClosureKwargs
Owned MLX closure with keyword arguments, released exactly once.
MlxClosureValueAndGrad
Owned MLX value-and-gradient closure, released exactly once.
MlxCompileCache
Owned MLX compile cache handle released exactly once.
MlxCompiledDualHeadNorm
Owns the compiled dual-output query/key head-normalization graph.
MlxCompiledElementwiseGraphs
Retained shapeless compilations for elementwise model graph composites.
MlxCompiledGatedDeltaRateActivation
Owns the compiled dual-output rate-activation graph.
MlxCompiledMultiOutputGraph
Owns one compiled MLX graph that produces an ordered vector of outputs.
MlxCompiledSwiGlu
Owned shapeless MLX compilation of silu(gate) * input.
MlxCudaKernel
Owned MLX-C CUDA kernel handle.
MlxCudaKernelConfig
Owned MLX-C CUDA kernel configuration handle.
MlxDeviceHandle
Owned MLX device handle released exactly once through the official C API.
MlxDeviceInfo
Owned MLX device information table released exactly once.
MlxDistributedGroup
Owned MLX distributed group handle released exactly once.
MlxFunctionExporter
Owned MLX function exporter released exactly once.
MlxImportedFunction
Owned MLX imported function released exactly once.
MlxIoGguf
Owned MLX GGUF container handle released exactly once.
MlxIoReader
Owned MLX reader transport released exactly once.
MlxIoWriter
Owned MLX writer transport released exactly once.
MlxMapStringToArray
Owned MLX map from strings to arrays, released exactly once.
MlxMapStringToArrayIterator
Owned iterator over an MLX string-to-array map.
MlxMapStringToString
Owned MLX map from strings to strings, released exactly once.
MlxMapStringToStringIterator
Owned iterator over an MLX string-to-string map.
MlxMetalKernel
Owned MLX custom Metal kernel function compiled from a source string.
MlxMetalKernelConfig
MlxMetalKernelOutput
Output shape and dtype requested from one custom Metal kernel launch.
MlxNodeNamer
Owned MLX node namer released exactly once.
MlxStream
Owned MLX stream handle used to preserve runtime thread affinity.
MlxStreamThreadLocal
Per-thread MLX stream slot for one device.
MlxString
Owned MLX string handle released exactly once through the official C API.
MlxVectorInt
Owned MLX integer vector handle released exactly once through the C API.
MlxVectorStream
Owned MLX vector of streams; elements are copied out as owned stream handles.
MlxVectorString
Owned MLX vector of strings; elements are borrowed from the vector for each read and copied into Rust memory.
MlxVectorVectorArray
Owned MLX vector of array vectors.

Enums§

MlxDtype
MLX array element types exposed without leaking generated C declarations.
MlxMetalKernelTemplateArgument
Compile-time template argument passed to an MLX custom Metal kernel launch.

Functions§

active_memory_bytes
Samples process-local MLX active-memory bytes.
apply_metal_kernel_in_graph
Applies one custom Metal kernel inside a compiled-graph trace.
array_from_vector
Reads one array out of a closure’s input vector inside a graph builder.
cache_memory_bytes
Samples process-local MLX allocator-cache bytes.
clear_allocator_cache
Releases reclaimable MLX allocator-cache allocations.
clear_streams
Frees every stream created in this process except the defaults in use.
cuda_is_available
Whether the CUDA backend reports availability in this process.
disable_compile
Disables MLX’s compilation transform globally.
distributed_is_available
Whether MLX’s distributed backend is initialized in this process.
enable_compile
Re-enables MLX’s compilation transform globally.
export_function
Exports a closure and example arguments to an MLX function file.
export_function_with_keywords
Exports a kwargs closure and example arguments to an MLX function file.
export_graph_to_dot⚠
Writes the graph behind the outputs to a Graphviz DOT file.
graph_output_array
Runs one raw MLX operation inside a graph builder and owns its lazy output.
initialize_distributed
Initializes MLX’s distributed backend.
jvp_of_closure
Computes the Jacobian-vector product of a closure.
load_array_from_reader
Loads one array from an open reader transport.
load_safetensors
Loads the string-to-array and string-to-string safetensors maps from a file.
load_safetensors_from_reader
Loads the safetensors maps from an open reader transport.
memory_limit_bytes
Reads the allocator limit currently enforced by MLX.
metal_is_available
Whether the MLX-Metal backend is available in this process.
metallib_path
The linked MLX-Metal backend’s metallib path.
peak_memory_bytes
Samples process-local MLX peak active-memory bytes.
print_graph⚠
Prints the graph behind the outputs as text.
reset_peak_memory
Resets MLX’s process-local peak active-memory counter.
save_array_to_writer
Saves one array through an open writer transport.
save_safetensors
Saves the parameter and metadata maps to a safetensors file.
save_safetensors_to_writer
Saves the parameter and metadata maps through an open writer transport.
set_cache_limit
Applies the allocator-cache retention limit and returns the previous limit in bytes.
set_compile_mode
Selects the global compilation mode.
set_default_stream
Replaces the process’s default stream.
set_graph_output⚠
Publishes the single output of a single-output graph builder.
set_graph_output_vector⚠
Appends one lazy output to a multi-output graph’s result vector.
set_memory_limit
Applies the graph-evaluation active-memory guidance and returns the previous limit in bytes.
set_metallib_path
Overrides the metallib MLX loads.
set_wired_limit
Raises the wired-memory ceiling for MLX allocations.
synchronize
Waits for submitted work on the given stream to complete.
synchronize_all_defaults
Synchronizes work across all devices’ default streams.
value_and_grad_of_closure
Builds the value-and-gradient closure for the given closure.
vjp_of_closure
Computes the vector-Jacobian product of a closure.
vmap_replace
Replaces one traced array in a vector-of-arrays graph, the building block of the VMAP transform.
vmap_trace
Traces a closure for the VMAP transform, producing traced inputs and outputs plus the resolved axes.

Type Aliases§

MlxGraphBuilder
The MLX C closure ABI every compiled-graph builder follows: it receives the output vector to populate and the input vector to read, and returns zero on success.