mlx_rs/transforms/
mod.rs

1//! Function transforms
2//!
3//! This mod provides functions for automatic differentiation and other
4//! transformations on functions.
5//!
6//! **WARN**: Because function transforms including compilation works on
7//! the computation graph, the user must ensure that all `Array`s are passed
8//! as inputs to the function/closure. Closures with captured `Array`s may
9//! not work as expected and may lead to undefined behavior.
10//!
11//! # Automatic Differentiation
12//!
13//! Automatic differentiation in MLX works on functions rather than on implicit
14//! graphs.
15//!
16//! **NOTE**: If you are coming to MLX from PyTorch, you no longer need
17//! functions like backward, zero_grad, and detach, or properties like
18//! requires_grad.
19//!
20//! You can use the [`grad()`] and [`value_and_grad()`] function to compute
21//! gradients of more complex functions. These functions compute the gradient
22//! with respect to the first argument, in order to manually specify the the
23//! argument to compute the gradient with respect to, use
24//! [`grad_with_argnums()`] or [`value_and_grad_with_argnums()`].
25//!
26//! # Panics
27//!
28//! A panic in a Rust transform closure is caught before it reaches MLX. After MLX returns through
29//! the C ABI, the panic resumes in Rust with its original payload.
30//!
31//! ```rust
32//! use mlx_rs::{Array, error::Result, transforms::grad};
33//!
34//! fn f(x: &Array) -> Result<Array> {
35//!     x.square()
36//! }
37//!
38//! fn calculate_grad(func: impl Fn(&Array) -> Result<Array>, arg: &Array) -> Result<Array> {
39//!     grad(func)(arg)
40//! }
41//!
42//! let x = Array::from(1.5);
43//!
44//! let dfdx = calculate_grad(f, &x).unwrap();
45//! assert_eq!(dfdx.item_exact::<f32>(), 2.0 * 1.5);
46//!
47//! let dfdx2 = calculate_grad(|args| calculate_grad(f, args), &x).unwrap();
48//! assert_eq!(dfdx2.item_exact::<f32>(), 2.0);
49//! ```
50
51use mlx_sys::mlx_closure_value_and_grad;
52
53use crate::{
54    error::{get_and_clear_closure_error, Result},
55    module::ModuleParamRef,
56    utils::{guard::Guarded, Closure, VectorArray, SUCCESS},
57    Array,
58};
59
60pub mod compile;
61mod grad;
62mod keyed_value_and_grad;
63mod value_and_grad;
64
65pub use grad::*;
66pub use keyed_value_and_grad::*;
67pub use value_and_grad::*;
68
69/// Evaluate an iterator of [`Array`]s.
70pub fn eval<'a>(outputs: impl IntoIterator<Item = &'a Array>) -> Result<()> {
71    let vec = VectorArray::try_from_iter(outputs.into_iter())?;
72    <() as Guarded>::try_from_op(|_| unsafe { mlx_sys::mlx_eval(vec.as_ptr()) })
73}
74
75/// Evaluate a module's parameters.
76///
77/// This is a convenience function that flattens the parameters and evaluates them.
78pub fn eval_params(params: ModuleParamRef<'_>) -> Result<()> {
79    eval(params.flatten().values().copied())
80}
81
82/// Asynchronously evaluate an iterator of [`Array`]s.
83///
84/// Please note that this is not a rust async function.
85pub fn async_eval<'a>(outputs: impl IntoIterator<Item = &'a Array>) -> Result<()> {
86    let vec = VectorArray::try_from_iter(outputs.into_iter())?;
87    <() as Guarded>::try_from_op(|_| unsafe { mlx_sys::mlx_async_eval(vec.as_ptr()) })
88}
89
90/// Asynchronously evaluate a module's parameters.
91///
92/// This is a convenience function that flattens the parameters and evaluates them.
93pub fn async_eval_params(params: ModuleParamRef<'_>) -> Result<()> {
94    async_eval(params.flatten().values().copied())
95}
96
97#[inline]
98fn jvp_inner(
99    closure: Closure<'_>,
100    primals: &[Array],
101    tangents: &[Array],
102) -> Result<(Vec<Array>, Vec<Array>)> {
103    let c_primals = VectorArray::try_from_iter(primals.iter())?;
104    let c_tangents = VectorArray::try_from_iter(tangents.iter())?;
105
106    <(Vec<Array>, Vec<Array>) as Guarded>::try_from_op(|(res_0, res_1)| unsafe {
107        mlx_sys::mlx_jvp(
108            res_0,
109            res_1,
110            closure.as_ptr(),
111            c_primals.as_ptr(),
112            c_tangents.as_ptr(),
113        )
114    })
115    .map_err(|e| match get_and_clear_closure_error() {
116        Some(err) => err,
117        None => e,
118    })
119}
120
121/// Compute the Jacobian-vector product.
122///
123/// This computes the product of the Jacobian of a function `f` evaluated at
124/// `primals` with the `tangents`.
125///
126/// # Params:
127///
128/// - `f`: function which takes an array of `Array` and returns an array of
129///   `Array`
130/// - `primals`: array of `Array` at which to evaluate the Jacobian
131/// - `tangents`: array of `Array` which are the "vector" in the Jacobian-vector
132///   product.  The `tangents` should be the same in number, shape and type as
133///   the inputs of `f`, e.g. the `primals`
134///
135/// # Returns:
136///
137/// Array of the Jacobian-vector products which is the same in number, shape and
138/// type of the outputs of `f`
139pub fn jvp<'a, F>(f: F, primals: &[Array], tangents: &[Array]) -> Result<(Vec<Array>, Vec<Array>)>
140where
141    F: FnMut(&[Array]) -> Vec<Array> + 'a,
142{
143    let closure = Closure::new(f);
144    jvp_inner(closure, primals, tangents)
145}
146
147/// Similar to [`jvp`] but handles closures that can return an error.
148pub fn fallible_jvp<'a, F>(
149    f: F,
150    primals: &[Array],
151    tangents: &[Array],
152) -> Result<(Vec<Array>, Vec<Array>)>
153where
154    F: FnMut(&[Array]) -> Result<Vec<Array>> + 'a,
155{
156    let closure = Closure::new_fallible(f);
157    jvp_inner(closure, primals, tangents)
158}
159
160#[inline]
161fn vjp_inner(
162    closure: Closure<'_>,
163    primals: &[Array],
164    cotangents: &[Array],
165) -> Result<(Vec<Array>, Vec<Array>)> {
166    let c_primals = VectorArray::try_from_iter(primals.iter())?;
167    let c_cotangents = VectorArray::try_from_iter(cotangents.iter())?;
168
169    <(Vec<Array>, Vec<Array>) as Guarded>::try_from_op(|(res_0, res_1)| unsafe {
170        mlx_sys::mlx_vjp(
171            res_0,
172            res_1,
173            closure.as_ptr(),
174            c_primals.as_ptr(),
175            c_cotangents.as_ptr(),
176        )
177    })
178    .map_err(|e| match get_and_clear_closure_error() {
179        Some(err) => err,
180        None => e,
181    })
182}
183
184/// Compute the vector-Jacobian product.
185///
186/// Computes the product of the `cotangents` with the Jacobian of a function `f` evaluated at
187/// `primals`.
188///
189/// # Params:
190///
191/// - f: function which takes an array of `Array` and returns an array of `Array`
192/// - primals: array of `Array` at which to evaluate the Jacobian
193/// - cotangents: array of `Array` which are the "vector" in the vector-Jacobian product. The
194///   `cotangents` should be the same in number, shape and type as the outputs of `f`
195///
196/// # Returns:
197///
198/// array of the vector-Jacobian products which is the same in number, shape and type of the outputs
199/// of `f`
200pub fn vjp<'a, F>(f: F, primals: &[Array], cotangents: &[Array]) -> Result<(Vec<Array>, Vec<Array>)>
201where
202    F: FnMut(&[Array]) -> Vec<Array> + 'a,
203{
204    let closure = Closure::new(f);
205    vjp_inner(closure, primals, cotangents)
206}
207
208/// Similar to [`vjp`] but handles closures that can return an error.
209pub fn fallible_vjp<'a, F>(
210    f: F,
211    primals: &[Array],
212    cotangents: &[Array],
213) -> Result<(Vec<Array>, Vec<Array>)>
214where
215    F: FnMut(&[Array]) -> Result<Vec<Array>> + 'a,
216{
217    let closure = Closure::new_fallible(f);
218    vjp_inner(closure, primals, cotangents)
219}
220
221pub(crate) struct ClosureValueAndGrad {
222    pub(crate) c_closure_value_and_grad: mlx_closure_value_and_grad,
223}
224
225impl ClosureValueAndGrad {
226    pub fn as_ptr(&self) -> mlx_closure_value_and_grad {
227        self.c_closure_value_and_grad
228    }
229}
230
231impl Drop for ClosureValueAndGrad {
232    fn drop(&mut self) {
233        let status =
234            unsafe { mlx_sys::mlx_closure_value_and_grad_free(self.c_closure_value_and_grad) };
235        debug_assert_eq!(status, SUCCESS);
236    }
237}
238
239fn value_and_gradient(
240    value_and_grad: mlx_closure_value_and_grad,
241    arrays: impl Iterator<Item = impl AsRef<Array>>,
242) -> Result<(Vec<Array>, Vec<Array>)> {
243    let input_vector = VectorArray::try_from_iter(arrays)?;
244
245    <(Vec<Array>, Vec<Array>) as Guarded>::try_from_op(|(res_0, res_1)| unsafe {
246        mlx_sys::mlx_closure_value_and_grad_apply(
247            res_0,
248            res_1,
249            value_and_grad,
250            input_vector.as_ptr(),
251        )
252    })
253    .map_err(|e| match get_and_clear_closure_error() {
254        Some(err) => err,
255        None => e,
256    })
257}
258
259#[cfg(test)]
260mod tests {
261
262    use crate::{
263        array,
264        transforms::{jvp, vjp},
265        Array,
266    };
267
268    use super::*;
269
270    // The unit tests below are adapted from the mlx c++ codebase
271
272    #[test]
273    fn test_jvp() {
274        let f = |inputs: &[Array]| -> Vec<Array> { vec![&inputs[0] + &inputs[1]] };
275        let x = array!(1.0f32);
276        let y = array!(1.0f32);
277        let (out, dout) = jvp(f, &[x, y], &[array!(1.0f32), array!(3.0f32)]).unwrap();
278        assert_eq!(out[0].item_exact::<f32>(), 2.0f32);
279        assert_eq!(dout[0].item_exact::<f32>(), 4.0f32);
280    }
281
282    #[test]
283    fn test_jvp_with_error() {
284        let f = |inputs: &[Array]| -> Result<Vec<Array>> {
285            inputs[0].add(&inputs[1]).map(|res| vec![res])
286        };
287
288        // Success case
289        let x = array!(1.0f32);
290        let y = array!(1.0f32);
291        let (out, dout) = fallible_jvp(f, &[x, y], &[array!(1.0f32), array!(3.0f32)]).unwrap();
292        assert_eq!(out[0].item_exact::<f32>(), 2.0f32);
293        assert_eq!(dout[0].item_exact::<f32>(), 4.0f32);
294
295        // Error case
296        // Use non-broadcastable shapes
297        let a = array!([1.0, 2.0, 3.0]);
298        let b = array!([4.0, 5.0]);
299        let result = fallible_jvp(f, &[a, b], &[array!(1.0f32), array!(3.0f32)]);
300        assert!(result.is_err());
301
302        // Check that the error is not just "mlx_closure returned a non-zero value"
303        let err = result.unwrap_err();
304        assert!(!err.what().contains("non-zero value"))
305    }
306
307    #[test]
308    fn test_vjp() {
309        let f = |inputs: &[Array]| -> Vec<Array> { vec![&inputs[0] + &inputs[1]] };
310        let x = array!(1.0f32);
311        let y = array!(1.0f32);
312        let primals = vec![x, y];
313        let cotangents = vec![array!(1.0f32)];
314        let (out, dout) = vjp(f, &primals, &cotangents).unwrap();
315        assert_eq!(out[0].item_exact::<f32>(), 2.0f32);
316        assert_eq!(dout[0].item_exact::<f32>(), 1.0f32);
317    }
318
319    #[test]
320    fn test_vjp_with_error() {
321        let f = |inputs: &[Array]| -> Result<Vec<Array>> {
322            inputs[0].add(&inputs[1]).map(|res| vec![res])
323        };
324
325        // Success case
326        let x = array!(1.0f32);
327        let y = array!(1.0f32);
328        let primals = vec![x, y];
329        let cotangents = vec![array!(1.0f32)];
330        let (out, dout) = fallible_vjp(f, &primals, &cotangents).unwrap();
331        assert_eq!(out[0].item_exact::<f32>(), 2.0f32);
332        assert_eq!(dout[0].item_exact::<f32>(), 1.0f32);
333
334        // Error case
335        // Use non-broadcastable shapes
336        let a = array!([1.0, 2.0, 3.0]);
337        let b = array!([4.0, 5.0]);
338        let result = fallible_vjp(f, &[a, b], &[array!(1.0f32)]);
339        assert!(result.is_err());
340
341        // Check that the error is not just "mlx_closure returned a non-zero value"
342        let err = result.unwrap_err();
343        assert!(!err.what().contains("non-zero value"))
344    }
345}
mlx_rs/transforms/mod.rs

mlx_rs/transforms/
mod.rs
