Crate mlx_rs

Expand description

Unofficial rust bindings for the MLX framework.

§Table of Contents

Quick Start
Lazy Evaluation
Unified Memory
Indexing Arrays
Saving and Loading

§Threading

Array values may move between threads. MLX 0.32.2 keeps default streams per thread, so each worker must perform operations through its own default or thread-local scoped override. Stream and compiled closures are thread-affine and cannot move or be shared across threads. The scoped Rust override is thread-local, not asynchronous task-local, and does not propagate across .await.

§Quick Start

See also MLX python documentation

§Basics

use mlx_rs::{array, Dtype};

let a = array!([1, 2, 3, 4]);
assert_eq!(a.shape(), &[4]);
assert_eq!(a.dtype(), Dtype::Int32);

let b = array!([1.0, 2.0, 3.0, 4.0]);
assert_eq!(b.dtype(), Dtype::Float32);

Operations in MLX are lazy. Use Array::eval to evaluate the the output of an operation. Operations are also automatically evaluated when inspecting an array with Array::item_exact or Array::item_cast, printing an array, or attempting to obtain the underlying data with Array::as_slice.

use mlx_rs::{array, transforms::eval};

let a = array!([1, 2, 3, 4]);
let b = array!([1.0, 2.0, 3.0, 4.0]);

let c = &a + &b; // c is not evaluated
c.eval().unwrap(); // evaluates c

let d = &a + &b;
println!("{:?}", d); // evaluates d

let e = &a + &b;
let e_slice: &[f32] = e.as_slice(); // evaluates e

See Lazy Evaluation for more details.

§Function and Graph Transformations

TODO: https://github.com/oxiglade/mlx-rs/issues/214

TODO: also document that all Array in the args for function transformations

§Lazy Evaluation

See also MLX python documentation

§Why Lazy Evaluation

When you perform operations in MLX, no computation actually happens. Instead a compute graph is recorded. The actual computation only happens if an Array::eval is performed.

MLX uses lazy evaluation because it has some nice features, some of which we describe below.

§Transforming Compute Graphs

Lazy evaluation lets us record a compute graph without actually doing any computations. This is useful for function transformations like transforms::grad and graph optimizations.

Currently, MLX does not compile and rerun compute graphs. They are all generated dynamically. However, lazy evaluation makes it much easier to integrate compilation for future performance enhancements.

§Only Compute What You Use

In MLX you do not need to worry as much about computing outputs that are never used. For example:

fn fun(x: &Array) -> (Array, Array) {
    let a = cheap_fun(x);
    let b = expensive_fun(x);
    (a, b)
}

let (y, _) = fun(&x);

Here, we never actually compute the output of expensive_fun. Use this pattern with care though, as the graph of expensive_fun is still built, and that has some cost associated to it.

Similarly, lazy evaluation can be beneficial for saving memory while keeping code simple. Say you have a very large model Model implementing module::Module. You can instantiate this model with let model = Model::new(). Typically, this will initialize all of the weights as float32, but the initialization does not actually compute anything until you perform an eval(). If you update the model with float16 weights, your maximum consumed memory will be half that required if eager computation was used instead.

This pattern is simple to do in MLX thanks to lazy computation:

let mut model = Model::new();
model.load_safetensors("model.safetensors").unwrap();

§When to Evaluate

A common question is when to use eval(). The trade-off is between letting graphs get too large and not batching enough useful work.

For example

let mut a = array!([1, 2, 3, 4]);
let mut b = array!([1.0, 2.0, 3.0, 4.0]);

for _ in 0..100 {
    a = a + b;
    a.eval()?;
    b = b * 2.0;
    b.eval()?;
}

This is a bad idea because there is some fixed overhead with each graph evaluation. On the other hand, there is some slight overhead which grows with the compute graph size, so extremely large graphs (while computationally correct) can be costly.

Luckily, a wide range of compute graph sizes work pretty well with MLX: anything from a few tens of operations to many thousands of operations per evaluation should be okay.

Most numerical computations have an iterative outer loop (e.g. the iteration in stochastic gradient descent). A natural and usually efficient place to use eval() is at each iteration of this outer loop.

Here is a concrete example:

for batch in dataset {
    // Nothing has been evaluated yet
    let (loss, grad) = value_and_grad_fn(&mut model, batch)?;

    // Still nothing has been evaluated
    optimizer.update(&mut model, grad)?;

    // Evaluate the loss and the new parameters which will
    // run the full gradient computation and optimizer update
    eval_params(model.parameters())?;
}

An important behavior to be aware of is when the graph will be implicitly evaluated. Anytime you print an array, or otherwise access its memory via Array::as_slice, the graph will be evaluated. Saving arrays via Array::save_numpy or Array::save_safetensors (or any other MLX saving functions) will also evaluate the array.

Calling Array::item_exact or Array::item_cast on a scalar array will also evaluate it. In the example above, printing the loss (println!("{:?}", loss)) or pushing the loss scalar to a Vec (losses.push(loss.item_exact::<f32>())) would cause a graph evaluation. If these lines are before evaluating the loss and module parameters, then this will be a partial evaluation, computing only the forward pass.

Also, calling eval() on an array or set of arrays multiple times is perfectly fine. This is effectively a no-op.

Warning: Using scalar arrays for control-flow will cause an evaluation.

fn fun(x: &Array) -> Array {
    let (h, y) = first_layer(x);

    if y.gt(array!(0.5)).unwrap().item_exact() {
        second_layer_a(h)
    } else {
        second_layer_b(h)
    }
}

Using arrays for control flow should be done with care. The above example works and can even be used with gradient transformations. However, this can be very inefficient if evaluations are done too frequently.

§Unified Memory

See also MLX python documentation

Apple silicon has a unified memory architecture. The CPU and GPU have direct access to the same memory pool. MLX is designed to take advantage of that.

Concretely, when you make an array in MLX you don’t have to specify its location:

let a = mlx_rs::random::normal::<f32>(&[100], None, None, None).unwrap();
let b = mlx_rs::random::normal::<f32>(&[100], None, None, None).unwrap();

Both a and b live in unified memory.

In MLX, rather than moving arrays to devices, you specify the device when you run the operation. Any device can perform any operation on a and b without needing to move them from one memory location to another. For example:

mlx_rs::with_device(mlx_rs::Device::cpu(), || mlx_rs::ops::add(&a, &b)).unwrap();
mlx_rs::with_device(mlx_rs::Device::gpu(), || mlx_rs::ops::add(&a, &b)).unwrap();

In the above, both the CPU and the GPU will perform the same add operation.

TODO: The remaining python documentations states that the stream can be used to parallelize operations without worrying about racing conditions. We should check if this is true given that we’ve already observed data racing when executing unit tests in parallel.

§Indexing Arrays

See also MLX python documentation

Please refer to the indexing modules (ops::indexing) for more details.

§Saving and Loading

See also MLX python documentation

mlx-rs supports loading from .npy and .safetensors files and saving to .safetensors files. Module parameters and optimizer states can also be saved and loaded from .safetensors files.

type	load function	save function
`Array`	`Array::load_numpy`	`Array::save_numpy`
`HashMap<String, Array>`	`Array::load_safetensors`	`Array::save_safetensors`
`module::Module`	`module::ModuleParametersExt::load_safetensors`	`module::ModuleParametersExt::save_safetensors`
`optimizers::Optimizer`	`optimizers::OptimizerState::load_safetensors`	`optimizers::OptimizerState::save_safetensors`

§Function Transforms

See also MLX python documentation

Please refer to the transforms module (transforms) for more details.

§Compilation

See also MLX python documentation

Please refer to the compilation module (transforms::compile) for more details.

Modules§

builder: Defines helper traits for builder pattern
error: Custom error types and handler for the c ffi
fast: Fast implementations of commonly used multi-op functions.
fft: Fast Fourier Transform (FFT) and its inverse (IFFT) for one, two, and N dimensions.
io: GGUF container loading, inspection, construction, and saving.
linalg: Linear algebra operations.
losses: Loss functions
macros: Macros for mlx-rs.
memory: Process-global MLX allocator controls and observations.
metal: Metal-specific runtime configuration.
module: This mod defines the traits for neural network modules and parameters.
nested: Implements a nested hashmap
nn: Neural network support for MLX
ops: Operations
optimizers: Trait and implementations for optimizers.
quantization: Traits for quantization
random: Collection of functions related to random number generation
transforms: Function transforms
utils: Utility functions and types.

Macros§

absDeprecated: Macro generated for the function crate::ops::abs. See the function documentation for more details.
acosDeprecated: Macro generated for the function crate::ops::acos. See the function documentation for more details.
acoshDeprecated: Macro generated for the function crate::ops::acosh. See the function documentation for more details.
addDeprecated: Macro generated for the function crate::ops::add. See the function documentation for more details.
addmmDeprecated: Macro generated for the function crate::ops::addmm. See the function documentation for more details.
allDeprecated: Macro generated for the function crate::ops::all. See the function documentation for more details.
all_axesDeprecated: Macro generated for the function crate::ops::all_axes. See the function documentation for more details.
all_axisDeprecated: Macro generated for the function crate::ops::all_axis. See the function documentation for more details.
all_closeDeprecated: Macro generated for the function crate::ops::all_close. See the function documentation for more details.
anyDeprecated: Macro generated for the function crate::ops::any. See the function documentation for more details.
any_axesDeprecated: Macro generated for the function crate::ops::any_axes. See the function documentation for more details.
any_axisDeprecated: Macro generated for the function crate::ops::any_axis. See the function documentation for more details.
arangeDeprecated: Macro generated for the function crate::ops::arange. See the function documentation for more details.
argmaxDeprecated: Macro generated for the function crate::ops::indexing::argmax. See the function documentation for more details.
argmax_axisDeprecated: Macro generated for the function crate::ops::indexing::argmax_axis. See the function documentation for more details.
argminDeprecated: Macro generated for the function crate::ops::indexing::argmin. See the function documentation for more details.
argmin_axisDeprecated: Macro generated for the function crate::ops::indexing::argmin_axis. See the function documentation for more details.
argpartitionDeprecated: Macro generated for the function crate::ops::argpartition. See the function documentation for more details.
argpartition_axisDeprecated: Macro generated for the function crate::ops::argpartition_axis. See the function documentation for more details.
argsortDeprecated: Macro generated for the function crate::ops::argsort. See the function documentation for more details.
argsort_axisDeprecated: 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_eqDeprecated: Macro generated for the function crate::ops::array_eq. See the function documentation for more details.
as_stridedDeprecated: Macro generated for the function crate::ops::as_strided. See the function documentation for more details.
asinDeprecated: Macro generated for the function crate::ops::asin. See the function documentation for more details.
asinhDeprecated: 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_1dDeprecated: Macro generated for the function crate::ops::at_least_1d. See the function documentation for more details.
at_least_2dDeprecated: Macro generated for the function crate::ops::at_least_2d. See the function documentation for more details.
at_least_3dDeprecated: Macro generated for the function crate::ops::at_least_3d. See the function documentation for more details.
atanDeprecated: Macro generated for the function crate::ops::atan. See the function documentation for more details.
atan2Deprecated: Macro generated for the function crate::ops::atan2. See the function documentation for more details.
atanhDeprecated: Macro generated for the function crate::ops::atanh. See the function documentation for more details.
bernoulliDeprecated: Macro generated for the function crate::random::bernoulli. See the function documentation for more details.
block_masked_mmDeprecated: Macro generated for the function crate::ops::block_masked_mm. See the function documentation for more details.
broadcast_arraysDeprecated: Macro generated for the function crate::ops::broadcast_arrays. See the function documentation for more details.
broadcast_toDeprecated: Macro generated for the function crate::ops::broadcast_to. See the function documentation for more details.
categoricalDeprecated: Macro generated for the function crate::random::categorical. See the function documentation for more details.
ceilDeprecated: Macro generated for the function crate::ops::ceil. See the function documentation for more details.
choleskyDeprecated: Macro generated for the function crate::linalg::cholesky. See the function documentation for more details.
cholesky_invDeprecated: Macro generated for the function crate::linalg::cholesky_inv. See the function documentation for more details.
clipDeprecated: Macro generated for the function crate::ops::clip. See the function documentation for more details.
concatenateDeprecated: Macro generated for the function crate::ops::concatenate. See the function documentation for more details.
concatenate_axisDeprecated: Macro generated for the function crate::ops::concatenate_axis. See the function documentation for more details.
conv1dDeprecated: Macro generated for the function crate::ops::conv1d. See the function documentation for more details.
conv2dDeprecated: Macro generated for the function crate::ops::conv2d. See the function documentation for more details.
conv3dDeprecated: Macro generated for the function crate::ops::conv3d. See the function documentation for more details.
conv_generalDeprecated: Macro generated for the function crate::ops::conv_general. See the function documentation for more details.
conv_transpose1dDeprecated: Macro generated for the function crate::ops::conv_transpose1d. See the function documentation for more details.
conv_transpose2dDeprecated: Macro generated for the function crate::ops::conv_transpose2d. See the function documentation for more details.
conv_transpose3dDeprecated: Macro generated for the function crate::ops::conv_transpose3d. See the function documentation for more details.
cosDeprecated: Macro generated for the function crate::ops::cos. See the function documentation for more details.
coshDeprecated: Macro generated for the function crate::ops::cosh. See the function documentation for more details.
crossDeprecated: Macro generated for the function crate::linalg::cross. See the function documentation for more details.
cummaxDeprecated: Macro generated for the function crate::ops::cummax. See the function documentation for more details.
cumminDeprecated: Macro generated for the function crate::ops::cummin. See the function documentation for more details.
cumprodDeprecated: Macro generated for the function crate::ops::cumprod. See the function documentation for more details.
cumsumDeprecated: Macro generated for the function crate::ops::cumsum. See the function documentation for more details.
degreesDeprecated: Macro generated for the function crate::ops::degrees. See the function documentation for more details.
dequantizeDeprecated: Macro generated for the function crate::ops::dequantize. See the function documentation for more details.
diagDeprecated: Macro generated for the function crate::ops::diag. See the function documentation for more details.
diagonalDeprecated: Macro generated for the function crate::ops::diagonal. See the function documentation for more details.
divideDeprecated: Macro generated for the function crate::ops::divide. See the function documentation for more details.
divmodDeprecated: Macro generated for the function crate::ops::divmod. See the function documentation for more details.
eigDeprecated: Macro generated for the function crate::linalg::eig. See the function documentation for more details.
eighDeprecated: Macro generated for the function crate::linalg::eigh. See the function documentation for more details.
eigvalsDeprecated: Macro generated for the function crate::linalg::eigvals. See the function documentation for more details.
eigvalshDeprecated: Macro generated for the function crate::linalg::eigvalsh. See the function documentation for more details.
einsumDeprecated: Macro generated for the function crate::ops::einsum. See the function documentation for more details.
eqDeprecated: Macro generated for the function crate::ops::eq. See the function documentation for more details.
erfDeprecated: Macro generated for the function crate::ops::erf. See the function documentation for more details.
erfinvDeprecated: Macro generated for the function crate::ops::erfinv. See the function documentation for more details.
expDeprecated: Macro generated for the function crate::ops::exp. See the function documentation for more details.
expand_dimsDeprecated: Macro generated for the function crate::ops::expand_dims. See the function documentation for more details.
expand_dims_axesDeprecated: Macro generated for the function crate::ops::expand_dims_axes. See the function documentation for more details.
expm1Deprecated: Macro generated for the function crate::ops::expm1. See the function documentation for more details.
eyeDeprecated: Macro generated for the function crate::ops::eye. See the function documentation for more details.
fftDeprecated: Macro generated for the function crate::fft::fft. See the function documentation for more details.
fft2Deprecated: Macro generated for the function crate::fft::fft2. See the function documentation for more details.
fftnDeprecated: Macro generated for the function crate::fft::fftn. See the function documentation for more details.
fftshiftDeprecated: Macro generated for the function crate::fft::fftshift. See the function documentation for more details.
flattenDeprecated: Macro generated for the function crate::ops::flatten. See the function documentation for more details.
floorDeprecated: Macro generated for the function crate::ops::floor. See the function documentation for more details.
floor_divideDeprecated: Macro generated for the function crate::ops::floor_divide. See the function documentation for more details.
from_fp8Deprecated: Macro generated for the function crate::ops::from_fp8. See the function documentation for more details.
fullDeprecated: Macro generated for the function crate::ops::full. See the function documentation for more details.
full_likeDeprecated: Macro generated for the function crate::ops::full_like. See the function documentation for more details.
gather_mmDeprecated: Macro generated for the function crate::ops::gather_mm. See the function documentation for more details.
gather_qmmDeprecated: Macro generated for the function crate::ops::gather_qmm. See the function documentation for more details.
gather_singleDeprecated: Macro generated for the function crate::ops::indexing::gather_single. See the function documentation for more details.
geDeprecated: Macro generated for the function crate::ops::ge. See the function documentation for more details.
gtDeprecated: Macro generated for the function crate::ops::gt. See the function documentation for more details.
gumbelDeprecated: Macro generated for the function crate::random::gumbel. See the function documentation for more details.
identityDeprecated: Macro generated for the function crate::ops::identity. See the function documentation for more details.
ifftDeprecated: Macro generated for the function crate::fft::ifft. See the function documentation for more details.
ifft2Deprecated: Macro generated for the function crate::fft::ifft2. See the function documentation for more details.
ifftnDeprecated: Macro generated for the function crate::fft::ifftn. See the function documentation for more details.
ifftshiftDeprecated: Macro generated for the function crate::fft::ifftshift. See the function documentation for more details.
imagDeprecated: Macro generated for the function crate::ops::imag. See the function documentation for more details.
innerDeprecated: Macro generated for the function crate::ops::inner. See the function documentation for more details.
invDeprecated: Macro generated for the function crate::linalg::inv. See the function documentation for more details.
irfftDeprecated: Macro generated for the function crate::fft::irfft. See the function documentation for more details.
irfft2Deprecated: Macro generated for the function crate::fft::irfft2. See the function documentation for more details.
irfftnDeprecated: Macro generated for the function crate::fft::irfftn. See the function documentation for more details.
is_closeDeprecated: Macro generated for the function crate::ops::is_close. See the function documentation for more details.
is_infDeprecated: Macro generated for the function crate::ops::is_inf. See the function documentation for more details.
is_nanDeprecated: Macro generated for the function crate::ops::is_nan. See the function documentation for more details.
is_neg_infDeprecated: Macro generated for the function crate::ops::is_neg_inf. See the function documentation for more details.
is_pos_infDeprecated: Macro generated for the function crate::ops::is_pos_inf. See the function documentation for more details.
kronDeprecated: Macro generated for the function crate::ops::kron. See the function documentation for more details.
layer_normDeprecated: Macro generated for the function crate::fast::layer_norm. See the function documentation for more details.
leDeprecated: Macro generated for the function crate::ops::le. See the function documentation for more details.
linspaceDeprecated: Macro generated for the function crate::ops::linspace. See the function documentation for more details.
logDeprecated: Macro generated for the function crate::ops::log. See the function documentation for more details.
log2Deprecated: Macro generated for the function crate::ops::log2. See the function documentation for more details.
log1pDeprecated: Macro generated for the function crate::ops::log1p. See the function documentation for more details.
log10Deprecated: Macro generated for the function crate::ops::log10. See the function documentation for more details.
logaddexpDeprecated: Macro generated for the function crate::ops::logaddexp. See the function documentation for more details.
logical_andDeprecated: Macro generated for the function crate::ops::logical_and. See the function documentation for more details.
logical_notDeprecated: Macro generated for the function crate::ops::logical_not. See the function documentation for more details.
logical_orDeprecated: Macro generated for the function crate::ops::logical_or. See the function documentation for more details.
logsumexpDeprecated: Macro generated for the function crate::ops::logsumexp. See the function documentation for more details.
logsumexp_axesDeprecated: Macro generated for the function crate::ops::logsumexp_axes. See the function documentation for more details.
logsumexp_axisDeprecated: Macro generated for the function crate::ops::logsumexp_axis. See the function documentation for more details.
ltDeprecated: Macro generated for the function crate::ops::lt. See the function documentation for more details.
luDeprecated: Macro generated for the function crate::linalg::lu. See the function documentation for more details.
lu_factorDeprecated: Macro generated for the function crate::linalg::lu_factor. See the function documentation for more details.
masked_scatterDeprecated: Macro generated for the function crate::ops::indexing::masked_scatter. See the function documentation for more details.
matmulDeprecated: Macro generated for the function crate::ops::matmul. See the function documentation for more details.
maxDeprecated: Macro generated for the function crate::ops::max. See the function documentation for more details.
max_axesDeprecated: Macro generated for the function crate::ops::max_axes. See the function documentation for more details.
max_axisDeprecated: Macro generated for the function crate::ops::max_axis. See the function documentation for more details.
maximumDeprecated: Macro generated for the function crate::ops::maximum. See the function documentation for more details.
meanDeprecated: Macro generated for the function crate::ops::mean. See the function documentation for more details.
mean_axesDeprecated: Macro generated for the function crate::ops::mean_axes. See the function documentation for more details.
mean_axisDeprecated: Macro generated for the function crate::ops::mean_axis. See the function documentation for more details.
medianDeprecated: Macro generated for the function crate::ops::median. See the function documentation for more details.
median_axesDeprecated: Macro generated for the function crate::ops::median_axes. See the function documentation for more details.
median_axisDeprecated: Macro generated for the function crate::ops::median_axis. See the function documentation for more details.
minDeprecated: Macro generated for the function crate::ops::min. See the function documentation for more details.
min_axesDeprecated: Macro generated for the function crate::ops::min_axes. See the function documentation for more details.
min_axisDeprecated: Macro generated for the function crate::ops::min_axis. See the function documentation for more details.
minimumDeprecated: Macro generated for the function crate::ops::minimum. See the function documentation for more details.
move_axisDeprecated: Macro generated for the function crate::ops::move_axis. See the function documentation for more details.
multiplyDeprecated: Macro generated for the function crate::ops::multiply. See the function documentation for more details.
multivariate_normalDeprecated: Macro generated for the function crate::random::multivariate_normal. See the function documentation for more details.
neDeprecated: Macro generated for the function crate::ops::ne. See the function documentation for more details.
negativeDeprecated: Macro generated for the function crate::ops::negative. See the function documentation for more details.
normDeprecated: Macro generated for the function crate::linalg::norm. See the function documentation for more details.
norm_l2Deprecated: Macro generated for the function crate::linalg::norm_l2. See the function documentation for more details.
norm_matrixDeprecated: Macro generated for the function crate::linalg::norm_matrix. See the function documentation for more details.
normalDeprecated: Macro generated for the function crate::random::normal. See the function documentation for more details.
onesDeprecated: Macro generated for the function crate::ops::ones. See the function documentation for more details.
ones_dtypeDeprecated: Macro generated for the function crate::ops::ones_dtype. See the function documentation for more details.
ones_likeDeprecated: Macro generated for the function crate::ops::ones_like. See the function documentation for more details.
outerDeprecated: Macro generated for the function crate::ops::outer. See the function documentation for more details.
padDeprecated: Macro generated for the function crate::ops::pad. See the function documentation for more details.
partitionDeprecated: Macro generated for the function crate::ops::partition. See the function documentation for more details.
partition_axisDeprecated: Macro generated for the function crate::ops::partition_axis. See the function documentation for more details.
pinvDeprecated: Macro generated for the function crate::linalg::pinv. See the function documentation for more details.
powerDeprecated: Macro generated for the function crate::ops::power. See the function documentation for more details.
prodDeprecated: Macro generated for the function crate::ops::prod. See the function documentation for more details.
prod_axesDeprecated: Macro generated for the function crate::ops::prod_axes. See the function documentation for more details.
prod_axisDeprecated: Macro generated for the function crate::ops::prod_axis. See the function documentation for more details.
put_along_axisDeprecated: Macro generated for the function crate::ops::indexing::put_along_axis. See the function documentation for more details.
qrDeprecated: Macro generated for the function crate::linalg::qr. See the function documentation for more details.
quantizeDeprecated: Macro generated for the function crate::ops::quantize. See the function documentation for more details.
quantized_matmulDeprecated: Macro generated for the function crate::ops::quantized_matmul. See the function documentation for more details.
radiansDeprecated: Macro generated for the function crate::ops::radians. See the function documentation for more details.
randintDeprecated: Macro generated for the function crate::random::randint. See the function documentation for more details.
realDeprecated: Macro generated for the function crate::ops::real. See the function documentation for more details.
reciprocalDeprecated: Macro generated for the function crate::ops::reciprocal. See the function documentation for more details.
remainderDeprecated: Macro generated for the function crate::ops::remainder. See the function documentation for more details.
repeatDeprecated: Macro generated for the function crate::ops::repeat. See the function documentation for more details.
repeat_axisDeprecated: Macro generated for the function crate::ops::repeat_axis. See the function documentation for more details.
reshapeDeprecated: Macro generated for the function crate::ops::reshape. See the function documentation for more details.
rfftDeprecated: Macro generated for the function crate::fft::rfft. See the function documentation for more details.
rfft2Deprecated: Macro generated for the function crate::fft::rfft2. See the function documentation for more details.
rfftnDeprecated: Macro generated for the function crate::fft::rfftn. See the function documentation for more details.
rms_normDeprecated: Macro generated for the function crate::fast::rms_norm. See the function documentation for more details.
ropeDeprecated: Macro generated for the function crate::fast::rope. See the function documentation for more details.
rope_dynamicDeprecated: Macro generated for the function crate::fast::rope_dynamic. See the function documentation for more details.
roundDeprecated: Macro generated for the function crate::ops::round. See the function documentation for more details.
rsqrtDeprecated: Macro generated for the function crate::ops::rsqrt. See the function documentation for more details.
scaled_dot_product_attentionDeprecated: Macro generated for the function crate::fast::scaled_dot_product_attention. See the function documentation for more details.
scatter_add_singleDeprecated: Macro generated for the function crate::ops::indexing::scatter_add_single. See the function documentation for more details.
scatter_max_singleDeprecated: Macro generated for the function crate::ops::indexing::scatter_max_single. See the function documentation for more details.
scatter_min_singleDeprecated: Macro generated for the function crate::ops::indexing::scatter_min_single. See the function documentation for more details.
scatter_prod_singleDeprecated: Macro generated for the function crate::ops::indexing::scatter_prod_single. See the function documentation for more details.
scatter_singleDeprecated: Macro generated for the function crate::ops::indexing::scatter_single. See the function documentation for more details.
segmented_mmDeprecated: Macro generated for the function crate::ops::segmented_mm. See the function documentation for more details.
sigmoidDeprecated: Macro generated for the function crate::ops::sigmoid. See the function documentation for more details.
signDeprecated: Macro generated for the function crate::ops::sign. See the function documentation for more details.
sinDeprecated: Macro generated for the function crate::ops::sin. See the function documentation for more details.
sinhDeprecated: Macro generated for the function crate::ops::sinh. See the function documentation for more details.
softmaxDeprecated: Macro generated for the function crate::ops::softmax. See the function documentation for more details.
softmax_axesDeprecated: Macro generated for the function crate::ops::softmax_axes. See the function documentation for more details.
softmax_axisDeprecated: Macro generated for the function crate::ops::softmax_axis. See the function documentation for more details.
solveDeprecated: Macro generated for the function crate::linalg::solve. See the function documentation for more details.
solve_triangularDeprecated: Macro generated for the function crate::linalg::solve_triangular. See the function documentation for more details.
sortDeprecated: Macro generated for the function crate::ops::sort. See the function documentation for more details.
sort_axisDeprecated: Macro generated for the function crate::ops::sort_axis. See the function documentation for more details.
splitDeprecated: Macro generated for the function crate::ops::split. See the function documentation for more details.
split_sectionsDeprecated: Macro generated for the function crate::ops::split_sections. See the function documentation for more details.
sqrtDeprecated: Macro generated for the function crate::ops::sqrt. See the function documentation for more details.
squareDeprecated: Macro generated for the function crate::ops::square. See the function documentation for more details.
squeezeDeprecated: Macro generated for the function crate::ops::squeeze. See the function documentation for more details.
squeeze_axesDeprecated: Macro generated for the function crate::ops::squeeze_axes. See the function documentation for more details.
stackDeprecated: Macro generated for the function crate::ops::stack. See the function documentation for more details.
stack_axisDeprecated: Macro generated for the function crate::ops::stack_axis. See the function documentation for more details.
stdDeprecated: Macro generated for the function crate::ops::std. See the function documentation for more details.
std_axesDeprecated: Macro generated for the function crate::ops::std_axes. See the function documentation for more details.
std_axisDeprecated: Macro generated for the function crate::ops::std_axis. See the function documentation for more details.
subtractDeprecated: Macro generated for the function crate::ops::subtract. See the function documentation for more details.
sumDeprecated: Macro generated for the function crate::ops::sum. See the function documentation for more details.
sum_axesDeprecated: Macro generated for the function crate::ops::sum_axes. See the function documentation for more details.
sum_axisDeprecated: Macro generated for the function crate::ops::sum_axis. See the function documentation for more details.
svdDeprecated: Macro generated for the function crate::linalg::svd. See the function documentation for more details.
swap_axesDeprecated: Macro generated for the function crate::ops::swap_axes. See the function documentation for more details.
takeDeprecated: Macro generated for the function crate::ops::indexing::take. See the function documentation for more details.
take_along_axisDeprecated: Macro generated for the function crate::ops::indexing::take_along_axis. See the function documentation for more details.
take_axisDeprecated: Macro generated for the function crate::ops::indexing::take_axis. See the function documentation for more details.
tanDeprecated: Macro generated for the function crate::ops::tan. See the function documentation for more details.
tanhDeprecated: Macro generated for the function crate::ops::tanh. See the function documentation for more details.
tensordot_axesDeprecated: Macro generated for the function crate::ops::tensordot_axes. See the function documentation for more details.
tensordot_axisDeprecated: Macro generated for the function crate::ops::tensordot_axis. See the function documentation for more details.
tileDeprecated: Macro generated for the function crate::ops::tile. See the function documentation for more details.
to_fp8Deprecated: Macro generated for the function crate::ops::to_fp8. See the function documentation for more details.
topkDeprecated: Macro generated for the function crate::ops::indexing::topk. See the function documentation for more details.
topk_axisDeprecated: Macro generated for the function crate::ops::indexing::topk_axis. See the function documentation for more details.
transposeDeprecated: Macro generated for the function crate::ops::transpose. See the function documentation for more details.
transpose_axesDeprecated: Macro generated for the function crate::ops::transpose_axes. See the function documentation for more details.
triDeprecated: Macro generated for the function crate::ops::tri. See the function documentation for more details.
tri_invDeprecated: Macro generated for the function crate::linalg::tri_inv. See the function documentation for more details.
trilDeprecated: Macro generated for the function crate::ops::tril. See the function documentation for more details.
triuDeprecated: Macro generated for the function crate::ops::triu. See the function documentation for more details.
truncated_normalDeprecated: Macro generated for the function crate::random::truncated_normal. See the function documentation for more details.
unflattenDeprecated: Macro generated for the function crate::ops::unflatten. See the function documentation for more details.
uniformDeprecated: Macro generated for the function crate::random::uniform. See the function documentation for more details.
varDeprecated: Macro generated for the function crate::ops::var. See the function documentation for more details.
var_axesDeprecated: Macro generated for the function crate::ops::var_axes. See the function documentation for more details.
var_axisDeprecated: Macro generated for the function crate::ops::var_axis. See the function documentation for more details.
whichDeprecated: Macro generated for the function crate::ops::which. See the function documentation for more details.
zerosDeprecated: Macro generated for the function crate::ops::zeros. See the function documentation for more details.
zeros_dtypeDeprecated: Macro generated for the function crate::ops::zeros_dtype. See the function documentation for more details.
zeros_likeDeprecated: 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
Stream: A stream of evaluation attached to a particular device.
StreamOrDevice: Parameter type for all MLX operations.

Enums§

Axes: Axis selection for operations that accept all, one, or several axes.
DeviceType: Type of device.
Dtype: Array element type

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§

stop_gradient: Stop gradients from being computed.
stop_gradient_device: Stop gradients from being computed.
task_local_default_streamDeprecated: Gets the thread-local scoped default stream.
thread_local_default_stream: Gets the thread-local scoped default stream.
with_device: Uses the default stream on device for operations constructed during f.
with_new_default_streamDeprecated: Uses a given default stream for the duration of f.
with_stream: Uses stream for operations constructed during f.

Type Aliases§

complex64: Type alias for num_complex::Complex<f32>.

Crate mlx_rs

Crate mlx_rs Copy item path

§Table of Contents

§Threading

§Quick Start

§Basics

§Function and Graph Transformations

§Lazy Evaluation

§Why Lazy Evaluation

§Transforming Compute Graphs

§Only Compute What You Use

§When to Evaluate

§Unified Memory

§Indexing Arrays

§Saving and Loading

§Function Transforms

§Compilation

Modules§

Macros§

Structs§

Enums§

Traits§

Functions§

Type Aliases§

Crate mlx_rs