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

mlx_rs/transforms/
mod.rs
