Crate safemlx 

Safe, low-level Rust bindings to MLX.

safemlx exposes MLX arrays, devices, streams, native operators, function transforms, and accelerator/runtime facilities. It intentionally does not define neural-network layers, models, optimizers, checkpoint policy, or other framework abstractions; those belong to consumers such as eredu-backend-mlx.

Operations are lazy and are scheduled on an explicit Stream. Call Array::evaluated (or a function in transforms) to submit work before accessing host data. Native MLX .npy and safetensors I/O is exposed on Array; higher-level checkpoint formats are outside this crate.

Modules

distributed

Distributed communication groups and operations.

error

Custom error types and handler for the c ffi

fast

Bindings to MLX fast operators and custom-kernel handles.

fft

Fast Fourier Transform (FFT) and its inverse (IFFT) for one, two, and N dimensions.

linalg

Linear algebra operations.

macros

Macros for safemlx.

memory

MLX allocator memory statistics.

metal

Metal backend configuration.

ops

Operations

random

Collection of functions related to random number generation

system

Safe process and operating-system resource observations used by runtimes.

transforms

Function transforms

utils

Utility functions and types.

Macros

abs

Macro generated for the function crate::ops::abs. See the function documentation for more details.

acos

Macro generated for the function crate::ops::acos. See the function documentation for more details.

acosh

Macro generated for the function crate::ops::acosh. See the function documentation for more details.

add

Macro generated for the function crate::ops::add. See the function documentation for more details.

addmm

Macro generated for the function crate::ops::addmm. See the function documentation for more details.

all

Macro generated for the function crate::ops::all. See the function documentation for more details.

all_axes

Macro generated for the function crate::ops::all_axes. See the function documentation for more details.

all_axis

Macro generated for the function crate::ops::all_axis. See the function documentation for more details.

all_close

Macro generated for the function crate::ops::all_close. See the function documentation for more details.

any

Macro generated for the function crate::ops::any. See the function documentation for more details.

any_axes

Macro generated for the function crate::ops::any_axes. See the function documentation for more details.

any_axis

Macro generated for the function crate::ops::any_axis. See the function documentation for more details.

arange

Macro generated for the function crate::ops::arange. See the function documentation for more details.

argmax

Macro generated for the function crate::ops::indexing::argmax. See the function documentation for more details.

argmax_axis

Macro generated for the function crate::ops::indexing::argmax_axis. See the function documentation for more details.

argmin

Macro generated for the function crate::ops::indexing::argmin. See the function documentation for more details.

argmin_axis

Macro generated for the function crate::ops::indexing::argmin_axis. See the function documentation for more details.

argpartition

Macro generated for the function crate::ops::argpartition. See the function documentation for more details.

argpartition_axis

Macro generated for the function crate::ops::argpartition_axis. See the function documentation for more details.

argsort

Macro generated for the function crate::ops::argsort. See the function documentation for more details.

argsort_axis

Macro generated for the function crate::ops::argsort_axis. See the function documentation for more details.

array

A helper macro to create an array with up to 3 dimensions.

array_eq

Macro generated for the function crate::ops::array_eq. See the function documentation for more details.

as_strided

Macro generated for the function crate::ops::as_strided. See the function documentation for more details.

asin

Macro generated for the function crate::ops::asin. See the function documentation for more details.

asinh

Macro generated for the function crate::ops::asinh. See the function documentation for more details.

assert_array_eq

Asserts that two arrays are equal.

at_least_1d

Macro generated for the function crate::ops::at_least_1d. See the function documentation for more details.

at_least_2d

Macro generated for the function crate::ops::at_least_2d. See the function documentation for more details.

at_least_3d

Macro generated for the function crate::ops::at_least_3d. See the function documentation for more details.

atan

Macro generated for the function crate::ops::atan. See the function documentation for more details.

atan2

Macro generated for the function crate::ops::atan2. See the function documentation for more details.

atanh

Macro generated for the function crate::ops::atanh. See the function documentation for more details.

bernoulli

Macro generated for the function crate::random::bernoulli. See the function documentation for more details.

block_masked_mm

Macro generated for the function crate::ops::block_masked_mm. See the function documentation for more details.

broadcast_arrays

Macro generated for the function crate::ops::broadcast_arrays. See the function documentation for more details.

broadcast_to

Macro generated for the function crate::ops::broadcast_to. See the function documentation for more details.

categorical

Macro generated for the function crate::random::categorical. See the function documentation for more details.

ceil

Macro generated for the function crate::ops::ceil. See the function documentation for more details.

cholesky

Macro generated for the function crate::linalg::cholesky. See the function documentation for more details.

cholesky_inv

Macro generated for the function crate::linalg::cholesky_inv. See the function documentation for more details.

clip

Macro generated for the function crate::ops::clip. See the function documentation for more details.

concatenate

Macro generated for the function crate::ops::concatenate. See the function documentation for more details.

concatenate_axis

Macro generated for the function crate::ops::concatenate_axis. See the function documentation for more details.

contiguous

Macro generated for the function crate::ops::contiguous. See the function documentation for more details.

conv1d

Macro generated for the function crate::ops::conv1d. See the function documentation for more details.

conv2d

Macro generated for the function crate::ops::conv2d. See the function documentation for more details.

conv3d

Macro generated for the function crate::ops::conv3d. See the function documentation for more details.

conv_general

Macro generated for the function crate::ops::conv_general. See the function documentation for more details.

conv_transpose1d

Macro generated for the function crate::ops::conv_transpose1d. See the function documentation for more details.

conv_transpose2d

Macro generated for the function crate::ops::conv_transpose2d. See the function documentation for more details.

conv_transpose3d

Macro generated for the function crate::ops::conv_transpose3d. See the function documentation for more details.

copy

Macro generated for the function crate::ops::copy. See the function documentation for more details.

cos

Macro generated for the function crate::ops::cos. See the function documentation for more details.

cosh

Macro generated for the function crate::ops::cosh. See the function documentation for more details.

cross

Macro generated for the function crate::linalg::cross. See the function documentation for more details.

cummax

Macro generated for the function crate::ops::cummax. See the function documentation for more details.

cummin

Macro generated for the function crate::ops::cummin. See the function documentation for more details.

cumprod

Macro generated for the function crate::ops::cumprod. See the function documentation for more details.

cumsum

Macro generated for the function crate::ops::cumsum. See the function documentation for more details.

degrees

Macro generated for the function crate::ops::degrees. See the function documentation for more details.

dequantize

Macro generated for the function crate::ops::dequantize. See the function documentation for more details.

diag

Macro generated for the function crate::ops::diag. See the function documentation for more details.

diagonal

Macro generated for the function crate::ops::diagonal. See the function documentation for more details.

divide

Macro generated for the function crate::ops::divide. See the function documentation for more details.

divmod

Macro generated for the function crate::ops::divmod. See the function documentation for more details.

eig

Macro generated for the function crate::linalg::eig. See the function documentation for more details.

eigh

Macro generated for the function crate::linalg::eigh. See the function documentation for more details.

eigvals

Macro generated for the function crate::linalg::eigvals. See the function documentation for more details.

eigvalsh

Macro generated for the function crate::linalg::eigvalsh. See the function documentation for more details.

einsum

Macro generated for the function crate::ops::einsum. See the function documentation for more details.

eq

Macro generated for the function crate::ops::eq. See the function documentation for more details.

erf

Macro generated for the function crate::ops::erf. See the function documentation for more details.

erfinv

Macro generated for the function crate::ops::erfinv. See the function documentation for more details.

exp

Macro generated for the function crate::ops::exp. See the function documentation for more details.

expand_dims

Macro generated for the function crate::ops::expand_dims. See the function documentation for more details.

expand_dims_axes

Macro generated for the function crate::ops::expand_dims_axes. See the function documentation for more details.

expm1

Macro generated for the function crate::ops::expm1. See the function documentation for more details.

eye

Macro generated for the function crate::ops::eye. See the function documentation for more details.

fft

Macro generated for the function crate::fft::fft. See the function documentation for more details.

fft2

Macro generated for the function crate::fft::fft2. See the function documentation for more details.

fftn

Macro generated for the function crate::fft::fftn. See the function documentation for more details.

fftshift

Macro generated for the function crate::fft::fftshift. See the function documentation for more details.

flatten

Macro generated for the function crate::ops::flatten. See the function documentation for more details.

floor

Macro generated for the function crate::ops::floor. See the function documentation for more details.

floor_divide

Macro generated for the function crate::ops::floor_divide. See the function documentation for more details.

from_fp8

Macro generated for the function crate::ops::from_fp8. See the function documentation for more details.

full

Macro generated for the function crate::ops::full. See the function documentation for more details.

full_like

Macro generated for the function crate::ops::full_like. See the function documentation for more details.

gather_mm

Macro generated for the function crate::ops::gather_mm. See the function documentation for more details.

gather_qmm

Macro generated for the function crate::ops::gather_qmm. See the function documentation for more details.

gather_single

Macro generated for the function crate::ops::indexing::gather_single. See the function documentation for more details.

ge

Macro generated for the function crate::ops::ge. See the function documentation for more details.

gt

Macro generated for the function crate::ops::gt. See the function documentation for more details.

gumbel

Macro generated for the function crate::random::gumbel. See the function documentation for more details.

identity

Macro generated for the function crate::ops::identity. See the function documentation for more details.

ifft

Macro generated for the function crate::fft::ifft. See the function documentation for more details.

ifft2

Macro generated for the function crate::fft::ifft2. See the function documentation for more details.

ifftn

Macro generated for the function crate::fft::ifftn. See the function documentation for more details.

ifftshift

Macro generated for the function crate::fft::ifftshift. See the function documentation for more details.

imag

Macro generated for the function crate::ops::imag. See the function documentation for more details.

inner

Macro generated for the function crate::ops::inner. See the function documentation for more details.

inv

Macro generated for the function crate::linalg::inv. See the function documentation for more details.

irfft

Macro generated for the function crate::fft::irfft. See the function documentation for more details.

irfft2

Macro generated for the function crate::fft::irfft2. See the function documentation for more details.

irfftn

Macro generated for the function crate::fft::irfftn. See the function documentation for more details.

is_close

Macro generated for the function crate::ops::is_close. See the function documentation for more details.

is_inf

Macro generated for the function crate::ops::is_inf. See the function documentation for more details.

is_nan

Macro generated for the function crate::ops::is_nan. See the function documentation for more details.

is_neg_inf

Macro generated for the function crate::ops::is_neg_inf. See the function documentation for more details.

is_pos_inf

Macro generated for the function crate::ops::is_pos_inf. See the function documentation for more details.

kron

Macro generated for the function crate::ops::kron. See the function documentation for more details.

layer_norm

Macro generated for the function crate::fast::layer_norm. See the function documentation for more details.

le

Macro generated for the function crate::ops::le. See the function documentation for more details.

linspace

Macro generated for the function crate::ops::linspace. See the function documentation for more details.

log

Macro generated for the function crate::ops::log. See the function documentation for more details.

log2

Macro generated for the function crate::ops::log2. See the function documentation for more details.

log1p

Macro generated for the function crate::ops::log1p. See the function documentation for more details.

log10

Macro generated for the function crate::ops::log10. See the function documentation for more details.

logaddexp

Macro generated for the function crate::ops::logaddexp. See the function documentation for more details.

logical_and

Macro generated for the function crate::ops::logical_and. See the function documentation for more details.

logical_not

Macro generated for the function crate::ops::logical_not. See the function documentation for more details.

logical_or

Macro generated for the function crate::ops::logical_or. See the function documentation for more details.

logsumexp

Macro generated for the function crate::ops::logsumexp. See the function documentation for more details.

logsumexp_axes

Macro generated for the function crate::ops::logsumexp_axes. See the function documentation for more details.

logsumexp_axis

Macro generated for the function crate::ops::logsumexp_axis. See the function documentation for more details.

lt

Macro generated for the function crate::ops::lt. See the function documentation for more details.

lu

Macro generated for the function crate::linalg::lu. See the function documentation for more details.

lu_factor

Macro generated for the function crate::linalg::lu_factor. See the function documentation for more details.

masked_scatter

Macro generated for the function crate::ops::indexing::masked_scatter. See the function documentation for more details.

matmul

Macro generated for the function crate::ops::matmul. See the function documentation for more details.

max

Macro generated for the function crate::ops::max. See the function documentation for more details.

max_axes

Macro generated for the function crate::ops::max_axes. See the function documentation for more details.

max_axis

Macro generated for the function crate::ops::max_axis. See the function documentation for more details.

maximum

Macro generated for the function crate::ops::maximum. See the function documentation for more details.

mean

Macro generated for the function crate::ops::mean. See the function documentation for more details.

mean_axes

Macro generated for the function crate::ops::mean_axes. See the function documentation for more details.

mean_axis

Macro generated for the function crate::ops::mean_axis. See the function documentation for more details.

median

Macro generated for the function crate::ops::median. See the function documentation for more details.

median_axes

Macro generated for the function crate::ops::median_axes. See the function documentation for more details.

median_axis

Macro generated for the function crate::ops::median_axis. See the function documentation for more details.

min

Macro generated for the function crate::ops::min. See the function documentation for more details.

min_axes

Macro generated for the function crate::ops::min_axes. See the function documentation for more details.

min_axis

Macro generated for the function crate::ops::min_axis. See the function documentation for more details.

minimum

Macro generated for the function crate::ops::minimum. See the function documentation for more details.

move_axis

Macro generated for the function crate::ops::move_axis. See the function documentation for more details.

multiply

Macro generated for the function crate::ops::multiply. See the function documentation for more details.

multivariate_normal

Macro generated for the function crate::random::multivariate_normal. See the function documentation for more details.

ne

Macro generated for the function crate::ops::ne. See the function documentation for more details.

negative

Macro generated for the function crate::ops::negative. See the function documentation for more details.

norm

Macro generated for the function crate::linalg::norm. See the function documentation for more details.

norm_l2

Macro generated for the function crate::linalg::norm_l2. See the function documentation for more details.

norm_matrix

Macro generated for the function crate::linalg::norm_matrix. See the function documentation for more details.

normal

Macro generated for the function crate::random::normal. See the function documentation for more details.

ones

Macro generated for the function crate::ops::ones. See the function documentation for more details.

ones_dtype

Macro generated for the function crate::ops::ones_dtype. See the function documentation for more details.

ones_like

Macro generated for the function crate::ops::ones_like. See the function documentation for more details.

outer

Macro generated for the function crate::ops::outer. See the function documentation for more details.

pad

Macro generated for the function crate::ops::pad. See the function documentation for more details.

partition

Macro generated for the function crate::ops::partition. See the function documentation for more details.

partition_axis

Macro generated for the function crate::ops::partition_axis. See the function documentation for more details.

pinv

Macro generated for the function crate::linalg::pinv. See the function documentation for more details.

power

Macro generated for the function crate::ops::power. See the function documentation for more details.

prod

Macro generated for the function crate::ops::prod. See the function documentation for more details.

prod_axes

Macro generated for the function crate::ops::prod_axes. See the function documentation for more details.

prod_axis

Macro generated for the function crate::ops::prod_axis. See the function documentation for more details.

put_along_axis

Macro generated for the function crate::ops::indexing::put_along_axis. See the function documentation for more details.

qr

Macro generated for the function crate::linalg::qr. See the function documentation for more details.

quantize

Macro generated for the function crate::ops::quantize. See the function documentation for more details.

quantized_matmul

Macro generated for the function crate::ops::quantized_matmul. See the function documentation for more details.

radians

Macro generated for the function crate::ops::radians. See the function documentation for more details.

randint

Macro generated for the function crate::random::randint. See the function documentation for more details.

real

Macro generated for the function crate::ops::real. See the function documentation for more details.

reciprocal

Macro generated for the function crate::ops::reciprocal. See the function documentation for more details.

remainder

Macro generated for the function crate::ops::remainder. See the function documentation for more details.

repeat

Macro generated for the function crate::ops::repeat. See the function documentation for more details.

repeat_axis

Macro generated for the function crate::ops::repeat_axis. See the function documentation for more details.

reshape

Macro generated for the function crate::ops::reshape. See the function documentation for more details.

rfft

Macro generated for the function crate::fft::rfft. See the function documentation for more details.

rfft2

Macro generated for the function crate::fft::rfft2. See the function documentation for more details.

rfftn

Macro generated for the function crate::fft::rfftn. See the function documentation for more details.

rms_norm

Macro generated for the function crate::fast::rms_norm. See the function documentation for more details.

rope

Macro generated for the function crate::fast::rope. See the function documentation for more details.

rope_dynamic

Macro generated for the function crate::fast::rope_dynamic. See the function documentation for more details.

round

Macro generated for the function crate::ops::round. See the function documentation for more details.

rsqrt

Macro generated for the function crate::ops::rsqrt. See the function documentation for more details.

scaled_dot_product_attention

Macro generated for the function crate::fast::scaled_dot_product_attention. See the function documentation for more details.

scatter

Macro generated for the function crate::ops::indexing::scatter. See the function documentation for more details.

scatter_add

Macro generated for the function crate::ops::indexing::scatter_add. See the function documentation for more details.

scatter_add_nd

Macro generated for the function crate::ops::indexing::scatter_add_nd. See the function documentation for more details.

scatter_add_single

Macro generated for the function crate::ops::indexing::scatter_add_single. See the function documentation for more details.

scatter_max_single

Macro generated for the function crate::ops::indexing::scatter_max_single. See the function documentation for more details.

scatter_min_single

Macro generated for the function crate::ops::indexing::scatter_min_single. See the function documentation for more details.

scatter_nd

Macro generated for the function crate::ops::indexing::scatter_nd. See the function documentation for more details.

scatter_prod_single

Macro generated for the function crate::ops::indexing::scatter_prod_single. See the function documentation for more details.

scatter_single

Macro generated for the function crate::ops::indexing::scatter_single. See the function documentation for more details.

segment_sum

Macro generated for the function crate::ops::segment_sum. See the function documentation for more details.

segmented_mm

Macro generated for the function crate::ops::segmented_mm. See the function documentation for more details.

sigmoid

Macro generated for the function crate::ops::sigmoid. See the function documentation for more details.

sign

Macro generated for the function crate::ops::sign. See the function documentation for more details.

sin

Macro generated for the function crate::ops::sin. See the function documentation for more details.

sinh

Macro generated for the function crate::ops::sinh. See the function documentation for more details.

softmax

Macro generated for the function crate::ops::softmax. See the function documentation for more details.

softmax_axes

Macro generated for the function crate::ops::softmax_axes. See the function documentation for more details.

softmax_axis

Macro generated for the function crate::ops::softmax_axis. See the function documentation for more details.

solve

Macro generated for the function crate::linalg::solve. See the function documentation for more details.

solve_triangular

Macro generated for the function crate::linalg::solve_triangular. See the function documentation for more details.

sort

Macro generated for the function crate::ops::sort. See the function documentation for more details.

sort_axis

Macro generated for the function crate::ops::sort_axis. See the function documentation for more details.

split

Macro generated for the function crate::ops::split. See the function documentation for more details.

split_sections

Macro generated for the function crate::ops::split_sections. See the function documentation for more details.

sqrt

Macro generated for the function crate::ops::sqrt. See the function documentation for more details.

square

Macro generated for the function crate::ops::square. See the function documentation for more details.

squeeze

Macro generated for the function crate::ops::squeeze. See the function documentation for more details.

squeeze_axes

Macro generated for the function crate::ops::squeeze_axes. See the function documentation for more details.

stack

Macro generated for the function crate::ops::stack. See the function documentation for more details.

stack_axis

Macro generated for the function crate::ops::stack_axis. See the function documentation for more details.

std

Macro generated for the function crate::ops::std. See the function documentation for more details.

std_axes

Macro generated for the function crate::ops::std_axes. See the function documentation for more details.

std_axis

Macro generated for the function crate::ops::std_axis. See the function documentation for more details.

subtract

Macro generated for the function crate::ops::subtract. See the function documentation for more details.

sum

Macro generated for the function crate::ops::sum. See the function documentation for more details.

sum_axes

Macro generated for the function crate::ops::sum_axes. See the function documentation for more details.

sum_axis

Macro generated for the function crate::ops::sum_axis. See the function documentation for more details.

svd

Macro generated for the function crate::linalg::svd. See the function documentation for more details.

swap_axes

Macro generated for the function crate::ops::swap_axes. See the function documentation for more details.

take

Macro generated for the function crate::ops::indexing::take. See the function documentation for more details.

take_along_axis

Macro generated for the function crate::ops::indexing::take_along_axis. See the function documentation for more details.

take_axis

Macro generated for the function crate::ops::indexing::take_axis. See the function documentation for more details.

tan

Macro generated for the function crate::ops::tan. See the function documentation for more details.

tanh

Macro generated for the function crate::ops::tanh. See the function documentation for more details.

tensordot_axes

Macro generated for the function crate::ops::tensordot_axes. See the function documentation for more details.

tensordot_axis

Macro generated for the function crate::ops::tensordot_axis. See the function documentation for more details.

tile

Macro generated for the function crate::ops::tile. See the function documentation for more details.

to_fp8

Macro generated for the function crate::ops::to_fp8. See the function documentation for more details.

topk

Macro generated for the function crate::ops::indexing::topk. See the function documentation for more details.

topk_axis

Macro generated for the function crate::ops::indexing::topk_axis. See the function documentation for more details.

transpose

Macro generated for the function crate::ops::transpose. See the function documentation for more details.

transpose_axes

Macro generated for the function crate::ops::transpose_axes. See the function documentation for more details.

tri

Macro generated for the function crate::ops::tri. See the function documentation for more details.

tri_inv

Macro generated for the function crate::linalg::tri_inv. See the function documentation for more details.

tril

Macro generated for the function crate::ops::tril. See the function documentation for more details.

triu

Macro generated for the function crate::ops::triu. See the function documentation for more details.

truncated_normal

Macro generated for the function crate::random::truncated_normal. See the function documentation for more details.

unflatten

Macro generated for the function crate::ops::unflatten. See the function documentation for more details.

uniform

Macro generated for the function crate::random::uniform. See the function documentation for more details.

var

Macro generated for the function crate::ops::var. See the function documentation for more details.

var_axes

Macro generated for the function crate::ops::var_axes. See the function documentation for more details.

var_axis

Macro generated for the function crate::ops::var_axis. See the function documentation for more details.

zeros

Macro generated for the function crate::ops::zeros. See the function documentation for more details.

zeros_dtype

Macro generated for the function crate::ops::zeros_dtype. See the function documentation for more details.

zeros_like

Macro generated for the function crate::ops::zeros_like. See the function documentation for more details.

Structs

Array

An n-dimensional array.

Device

Representation of a Device in MLX.

DtypeIter

An iterator over the variants of Dtype

EvaluatedArray

An evaluated array with materialized storage available for host reads.

Event

A single-shot completion event independent of the selected MLX device backend.

HostTransferBuffer

An owned, typed, host-addressable transfer allocation.

HostTransferMemoryStats

Process-wide physical allocations owned by host-transfer buffers of one kind.

ImmutableHostTransferBuffer

Immutable, shareable host-transfer storage.

PendingDeviceTransfer

A host-to-array transfer whose array cannot be accessed before completion.

PendingHostTransfer

An array-to-host transfer whose buffer cannot be accessed before completion.

Stream

A stream of evaluation attached to a particular device.

SubmittedDeviceTransfer

A submitted immutable-host-to-device transfer.

TimedEvaluation

An asynchronously submitted evaluation with execution-timeline timestamps.

Enums

DeviceType

Type of device.

Dtype

Array element type

EventBackend

MLX device backend which owns an Event.

HostTransferPolicy

Requested semantics for an MLX host transfer allocation.

HostTransferStorageKind

Physical storage selected for a host transfer allocation.

Traits

ArrayElement

A marker trait for array elements.

FromNested

A helper trait to construct Array from nested arrays or slices.

FromScalar

A helper trait to construct Array from scalar values.

Functions

host_transfer_capacity_upper_bound

Returns a conservative physical capacity to reserve before allocation.

host_transfer_memory_stats

Returns native physical-allocation telemetry for one storage kind.

reset_host_transfer_peak_memory

Resets one storage kind’s physical high-water marks to current occupancy.

stop_gradient

Stop gradients from being computed.

Type Aliases

complex64

Type alias for num_complex::Complex<f32>.