mlx_rs/transforms/compile/
compile_with_state.rs

1//! Compilation of functions with state.
2//!
3//! # Unit tests
4//!
5//! See `mlx-rs/mlx-tests/tests/test_compile.rs` for unit tests.
6
7// TODO: there's plenty boilerplate code here but it's not clear how to reduce it
8
9use std::{
10    cell::{Cell, RefCell},
11    collections::BTreeSet,
12    marker::PhantomData,
13    rc::Rc,
14};
15
16use crate::{
17    error::Exception,
18    transforms::compile::CompiledState,
19    utils::{StateSnapshot, Updatable},
20    Array,
21};
22
23use super::{Closure, Compiled, Guarded, StateLayout, VectorArray};
24
25/// Similar to [`crate::transforms::compile`] but allows for functions that take
26/// a mutable reference to a state `U`.
27pub fn compile_with_state<F, U, A, O, E>(
28    f: F,
29    shapeless: impl Into<Option<bool>>,
30) -> impl for<'a> FnMut(&mut U, F::Args<'a>) -> Result<O, Exception>
31where
32    F: CompileWithState<U, A, O, E> + 'static,
33    U: Updatable,
34{
35    let shapeless = shapeless.into().unwrap_or(false);
36    let mut compiled = f.compile(shapeless);
37    move |state, args| compiled.call_mut(state, args)
38}
39
40/// A trait for functions that can be compiled with state.
41///
42/// This trait is used to compile a function that takes a mutable reference to a state
43/// and some arguments and returns a result.
44///
45/// # Generic parameters
46///
47/// - `U`: The type of the state.
48/// - `A`: The type of the arguments.
49/// - `O`: The type of the output.
50/// - `E`: The type of the exception.
51pub trait CompileWithState<U, A, O, E> {
52    /// The type of the arguments that the returned closure takes.
53    ///
54    /// This is needed to relax the lifetime requirements of the returned
55    /// closure. Otherwise, the arguments to the returned closure would have to
56    /// live longer than the closure itself.
57    type Args<'a>;
58
59    /// Compile the function.
60    fn compile(
61        self,
62        shapeless: bool,
63    ) -> impl for<'args> CallMutWithState<U, Self::Args<'args>, O, E>;
64}
65
66impl<F, U> CompileWithState<U, &[Array], Vec<Array>, ()> for F
67where
68    F: FnMut(&mut U, &[Array]) -> Vec<Array> + 'static,
69    U: Updatable,
70{
71    type Args<'a> = &'a [Array];
72
73    fn compile(
74        self,
75        shapeless: bool,
76    ) -> impl for<'args> CallMutWithState<U, Self::Args<'args>, Vec<Array>, ()> {
77        let state = CompiledState::new(self, shapeless);
78        Compiled::<F, _> {
79            f_marker: PhantomData,
80            state,
81        }
82    }
83}
84
85impl<F, U> CompileWithState<U, &Array, Array, ()> for F
86where
87    F: FnMut(&mut U, &Array) -> Array + 'static,
88    U: Updatable,
89{
90    type Args<'a> = &'a Array;
91
92    fn compile(
93        mut self,
94        shapeless: bool,
95    ) -> impl for<'args> CallMutWithState<U, Self::Args<'args>, Array, ()> {
96        let f = move |state: &mut U, args: &[Array]| -> Vec<Array> {
97            let result = (self)(state, &args[0]);
98            vec![result]
99        };
100        let state = CompiledState::new(f, shapeless);
101        Compiled::<F, _> {
102            f_marker: PhantomData,
103            state,
104        }
105    }
106}
107
108impl<F, U> CompileWithState<U, (&Array, &Array), Array, ()> for F
109where
110    F: FnMut(&mut U, (&Array, &Array)) -> Array + 'static,
111    U: Updatable,
112{
113    type Args<'a> = (&'a Array, &'a Array);
114
115    fn compile(
116        mut self,
117        shapeless: bool,
118    ) -> impl for<'args> CallMutWithState<U, Self::Args<'args>, Array, ()> {
119        let f = move |state: &mut U, args: &[Array]| -> Vec<Array> {
120            let result = (self)(state, (&args[0], &args[1]));
121            vec![result]
122        };
123        let state = CompiledState::new(f, shapeless);
124        Compiled::<F, _> {
125            f_marker: PhantomData,
126            state,
127        }
128    }
129}
130
131impl<F, U> CompileWithState<U, (&Array, &Array, &Array), Array, ()> for F
132where
133    F: FnMut(&mut U, (&Array, &Array, &Array)) -> Array + 'static,
134    U: Updatable,
135{
136    type Args<'a> = (&'a Array, &'a Array, &'a Array);
137
138    fn compile(
139        mut self,
140        shapeless: bool,
141    ) -> impl for<'args> CallMutWithState<U, Self::Args<'args>, Array, ()> {
142        let f = move |state: &mut U, args: &[Array]| -> Vec<Array> {
143            let result = (self)(state, (&args[0], &args[1], &args[2]));
144            vec![result]
145        };
146        let state = CompiledState::new(f, shapeless);
147        Compiled::<F, _> {
148            f_marker: PhantomData,
149            state,
150        }
151    }
152}
153
154impl<F, U> CompileWithState<U, &[Array], Vec<Array>, Exception> for F
155where
156    F: FnMut(&mut U, &[Array]) -> Result<Vec<Array>, Exception> + 'static,
157    U: Updatable,
158{
159    type Args<'a> = &'a [Array];
160
161    fn compile(
162        self,
163        shapeless: bool,
164    ) -> impl for<'args> CallMutWithState<U, Self::Args<'args>, Vec<Array>, Exception> {
165        let state = CompiledState::new(self, shapeless);
166        Compiled::<F, _> {
167            f_marker: PhantomData,
168            state,
169        }
170    }
171}
172
173impl<F, U> CompileWithState<U, &Array, Array, Exception> for F
174where
175    F: FnMut(&mut U, &Array) -> Result<Array, Exception> + 'static,
176    U: Updatable,
177{
178    type Args<'a> = &'a Array;
179
180    fn compile(
181        mut self,
182        shapeless: bool,
183    ) -> impl for<'args> CallMutWithState<U, Self::Args<'args>, Array, Exception> {
184        let f = move |state: &mut U, args: &[Array]| -> Result<Vec<Array>, Exception> {
185            let result = (self)(state, &args[0])?;
186            Ok(vec![result])
187        };
188        let state = CompiledState::new(f, shapeless);
189        Compiled::<F, _> {
190            f_marker: PhantomData,
191            state,
192        }
193    }
194}
195
196impl<F, U> CompileWithState<U, (&Array, &Array), Array, Exception> for F
197where
198    F: FnMut(&mut U, (&Array, &Array)) -> Result<Array, Exception> + 'static,
199    U: Updatable,
200{
201    type Args<'a> = (&'a Array, &'a Array);
202
203    fn compile(
204        mut self,
205        shapeless: bool,
206    ) -> impl for<'args> CallMutWithState<U, Self::Args<'args>, Array, Exception> {
207        let f = move |state: &mut U, args: &[Array]| -> Result<Vec<Array>, Exception> {
208            let result = (self)(state, (&args[0], &args[1]))?;
209            Ok(vec![result])
210        };
211        let state = CompiledState::new(f, shapeless);
212        Compiled::<F, _> {
213            f_marker: PhantomData,
214            state,
215        }
216    }
217}
218
219impl<F, U> CompileWithState<U, (&Array, &Array, &Array), Array, Exception> for F
220where
221    F: FnMut(&mut U, (&Array, &Array, &Array)) -> Result<Array, Exception> + 'static,
222    U: Updatable,
223{
224    type Args<'a> = (&'a Array, &'a Array, &'a Array);
225
226    fn compile(
227        mut self,
228        shapeless: bool,
229    ) -> impl for<'args> CallMutWithState<U, Self::Args<'args>, Array, Exception> {
230        let f = move |state: &mut U, args: &[Array]| -> Result<Vec<Array>, Exception> {
231            let result = (self)(state, (&args[0], &args[1], &args[2]))?;
232            Ok(vec![result])
233        };
234        let state = CompiledState::new(f, shapeless);
235        Compiled::<F, _> {
236            f_marker: PhantomData,
237            state,
238        }
239    }
240}
241
242/// A trait for functions that can be called with state.
243pub trait CallMutWithState<U, A, O, E> {
244    /// Call the function with the given state and arguments.
245    fn call_mut(&mut self, state: &mut U, args: A) -> Result<O, Exception>;
246}
247
248impl<U, F, G> CallMutWithState<U, &[Array], Vec<Array>, ()> for Compiled<F, G>
249where
250    F: FnMut(&mut U, &[Array]) -> Vec<Array>,
251    G: FnMut(&mut U, &[Array]) -> Vec<Array>,
252    U: Updatable,
253{
254    fn call_mut(&mut self, state: &mut U, args: &[Array]) -> Result<Vec<Array>, Exception> {
255        self.state.call_mut_with_state(state, args)
256    }
257}
258
259impl<U, F, G> CallMutWithState<U, &Array, Array, ()> for Compiled<F, G>
260where
261    F: FnMut(&mut U, &Array) -> Array,
262    G: FnMut(&mut U, &[Array]) -> Vec<Array>,
263    U: Updatable,
264{
265    fn call_mut(&mut self, state: &mut U, args: &Array) -> Result<Array, Exception> {
266        let args = std::slice::from_ref(args);
267        let result = self.state.call_mut_with_state(state, args)?;
268        Ok(result.into_iter().next().unwrap())
269    }
270}
271
272impl<U, F, G> CallMutWithState<U, (&Array, &Array), Array, ()> for Compiled<F, G>
273where
274    F: FnMut(&mut U, (&Array, &Array)) -> Array,
275    G: FnMut(&mut U, &[Array]) -> Vec<Array>,
276    U: Updatable,
277{
278    fn call_mut(&mut self, state: &mut U, args: (&Array, &Array)) -> Result<Array, Exception> {
279        let args = &[args.0, args.1];
280        let result = self.state.call_mut_with_state(state, args)?;
281        Ok(result.into_iter().next().unwrap())
282    }
283}
284
285impl<U, F, G> CallMutWithState<U, (&Array, &Array, &Array), Array, ()> for Compiled<F, G>
286where
287    F: FnMut(&mut U, (&Array, &Array, &Array)) -> Array,
288    G: FnMut(&mut U, &[Array]) -> Vec<Array>,
289    U: Updatable,
290{
291    fn call_mut(
292        &mut self,
293        state: &mut U,
294        args: (&Array, &Array, &Array),
295    ) -> Result<Array, Exception> {
296        let args = &[args.0, args.1, args.2];
297        let result = self.state.call_mut_with_state(state, args)?;
298        Ok(result.into_iter().next().unwrap())
299    }
300}
301
302impl<U, F, G> CallMutWithState<U, &[Array], Vec<Array>, Exception> for Compiled<F, G>
303where
304    F: FnMut(&mut U, &[Array]) -> Result<Vec<Array>, Exception>,
305    G: FnMut(&mut U, &[Array]) -> Result<Vec<Array>, Exception>,
306    U: Updatable,
307{
308    fn call_mut(&mut self, state: &mut U, args: &[Array]) -> Result<Vec<Array>, Exception> {
309        self.state.fallible_call_mut_with_state(state, args)
310    }
311}
312
313impl<U, F, G> CallMutWithState<U, &Array, Array, Exception> for Compiled<F, G>
314where
315    F: FnMut(&mut U, &Array) -> Result<Array, Exception>,
316    G: FnMut(&mut U, &[Array]) -> Result<Vec<Array>, Exception>,
317    U: Updatable,
318{
319    fn call_mut(&mut self, state: &mut U, args: &Array) -> Result<Array, Exception> {
320        let args = std::slice::from_ref(args);
321        let result = self.state.fallible_call_mut_with_state(state, args)?;
322        Ok(result.into_iter().next().unwrap())
323    }
324}
325
326impl<U, F, G> CallMutWithState<U, (&Array, &Array), Array, Exception> for Compiled<F, G>
327where
328    F: FnMut(&mut U, (&Array, &Array)) -> Result<Array, Exception>,
329    G: FnMut(&mut U, &[Array]) -> Result<Vec<Array>, Exception>,
330    U: Updatable,
331{
332    fn call_mut(&mut self, state: &mut U, args: (&Array, &Array)) -> Result<Array, Exception> {
333        let args = &[args.0, args.1];
334        let result = self.state.fallible_call_mut_with_state(state, args)?;
335        Ok(result.into_iter().next().unwrap())
336    }
337}
338
339impl<U, F, G> CallMutWithState<U, (&Array, &Array, &Array), Array, Exception> for Compiled<F, G>
340where
341    F: FnMut(&mut U, (&Array, &Array, &Array)) -> Result<Array, Exception>,
342    G: FnMut(&mut U, &[Array]) -> Result<Vec<Array>, Exception>,
343    U: Updatable,
344{
345    fn call_mut(
346        &mut self,
347        state: &mut U,
348        args: (&Array, &Array, &Array),
349    ) -> Result<Array, Exception> {
350        let args = &[args.0, args.1, args.2];
351        let result = self.state.fallible_call_mut_with_state(state, args)?;
352        Ok(result.into_iter().next().unwrap())
353    }
354}
355
356#[inline]
357fn call_mut_with_state_inner<U>(
358    inner_closure: Closure,
359    fun_id: usize,
360    shapeless: bool,
361    state: Rc<RefCell<&mut U>>,
362    args: &[impl AsRef<Array>],
363    num_function_outputs: Rc<Cell<Option<usize>>>,
364    state_layout: Rc<RefCell<Option<StateLayout>>>,
365) -> crate::error::Result<Vec<Array>>
366where
367    U: Updatable,
368{
369    // note: this will use the cached compile (via the id)
370    // but will be able to re-evaluate with fresh state if needed
371    let compiled = Closure::try_from_op(|res| unsafe {
372        let constants = &[];
373        mlx_sys::mlx_detail_compile(
374            res,
375            inner_closure.as_ptr(),
376            fun_id,
377            shapeless,
378            constants.as_ptr(),
379            0,
380        )
381    })?;
382
383    let inner_inputs_vector = {
384        let mut borrow = state.borrow_mut();
385        let state_inputs = borrow.state_projection()?.snapshot();
386        VectorArray::try_from_iter(
387            args.iter()
388                .map(AsRef::as_ref)
389                .chain(state_inputs.present_values()),
390        )?
391    };
392
393    // will compile the function (if needed) and evaluate the
394    // compiled graph
395    let result_vector = VectorArray::try_from_op(|res| unsafe {
396        mlx_sys::mlx_closure_apply(res, compiled.as_ptr(), inner_inputs_vector.as_ptr())
397    })?;
398
399    let result_plus_state_output: Vec<Array> = result_vector.try_into_values()?;
400
401    // The combined output layout is: [function_outputs..., state_arrays...]
402    // We captured the function output count during tracing to know where to split.
403    let num_fn_outputs = num_function_outputs.get().ok_or_else(|| {
404        Exception::custom(
405            "compile_with_state: internal error - function output count not captured during tracing"
406        )
407    })?;
408    let expected_state_layout = state_layout.borrow().clone().ok_or_else(|| {
409        Exception::custom(
410            "compile_with_state: internal error - state layout not captured during tracing",
411        )
412    })?;
413    validate_state_layout(
414        &expected_state_layout,
415        &mut **state.borrow_mut(),
416        "apply input",
417    )?;
418
419    let expected_state_outputs = expected_state_layout
420        .iter()
421        .filter(|entry| entry.is_present())
422        .count();
423    let expected_output_count = num_fn_outputs + expected_state_outputs;
424    if result_plus_state_output.len() != expected_output_count {
425        return Err(Exception::custom(format!(
426            "compile_with_state: invalid output count - expected {num_fn_outputs} function \
427             outputs and {} state outputs, got {} total outputs",
428            expected_state_outputs,
429            result_plus_state_output.len()
430        )));
431    }
432
433    let function_results = &result_plus_state_output[..num_fn_outputs];
434    let state_outputs = &result_plus_state_output[num_fn_outputs..];
435
436    let state_output =
437        StateSnapshot::from_layout_and_values(&expected_state_layout, state_outputs)?;
438    let output_layout = state_output.layout();
439    if output_layout != expected_state_layout {
440        return Err(Exception::custom(format!(
441            "compile_with_state: state output layout changed: expected \
442             {expected_state_layout:?}, got {output_layout:?}"
443        )));
444    }
445
446    state
447        .borrow_mut()
448        .state_projection()?
449        .restore(state_output, false)?;
450
451    // Return only the function results (not the state arrays)
452    Ok(function_results.to_vec())
453}
454
455fn current_state_layout(state: &mut impl Updatable) -> Result<StateLayout, Exception> {
456    Ok(state.state_projection()?.layout())
457}
458
459fn validate_state_layout(
460    expected: &[crate::utils::StateLayoutEntry],
461    state: &mut impl Updatable,
462    phase: &str,
463) -> Result<(), Exception> {
464    let actual = current_state_layout(state)?;
465    if actual == expected {
466        Ok(())
467    } else {
468        Err(Exception::custom(format!(
469            "compile_with_state: state layout changed at {phase}: expected {expected:?}, got \
470             {actual:?}"
471        )))
472    }
473}
474
475fn state_snapshot(state: &mut impl Updatable) -> Result<StateSnapshot, Exception> {
476    Ok(state.state_projection()?.snapshot())
477}
478
479fn restore_state(
480    state: &mut impl Updatable,
481    snapshot: &StateSnapshot,
482    reset_new: bool,
483) -> Result<(), Exception> {
484    Ok(state
485        .state_projection()?
486        .restore(snapshot.clone(), reset_new)?)
487}
488
489fn state_grew(before: &StateSnapshot, after: &StateSnapshot) -> bool {
490    let before = before
491        .iter()
492        .map(|(key, _)| key.to_owned())
493        .collect::<BTreeSet<_>>();
494    let after = after
495        .iter()
496        .map(|(key, _)| key.to_owned())
497        .collect::<BTreeSet<_>>();
498    before.len() < after.len() && before.is_subset(&after)
499}
500
501impl<F> CompiledState<F> {
502    fn call_mut_with_state<U>(
503        &mut self,
504        state: &mut U,
505        args: &[impl AsRef<Array>],
506    ) -> Result<Vec<Array>, Exception>
507    where
508        F: FnMut(&mut U, &[Array]) -> Vec<Array>,
509        U: Updatable,
510    {
511        self.call_mut_with_state_attempt(state, args, true)
512    }
513
514    fn call_mut_with_state_attempt<U>(
515        &mut self,
516        state: &mut U,
517        args: &[impl AsRef<Array>],
518        allow_growth_recovery: bool,
519    ) -> Result<Vec<Array>, Exception>
520    where
521        F: FnMut(&mut U, &[Array]) -> Vec<Array>,
522        U: Updatable,
523    {
524        let was_untraced = self.state_layout.is_none();
525        if let Some(expected) = self.state_layout.as_deref() {
526            validate_state_layout(expected, state, "call input")?;
527        }
528        let args_len = args.len();
529        let saved_state = state_snapshot(state)?;
530        let state = Rc::new(RefCell::new(state));
531        let f = &mut self.f;
532
533        // Cell to capture the number of function outputs during tracing
534        let num_function_outputs = Rc::new(Cell::new(self.num_function_outputs));
535        let num_fn_outputs_clone = Rc::clone(&num_function_outputs);
536        let state_layout = Rc::new(RefCell::new(self.state_layout.clone()));
537        let state_layout_clone = Rc::clone(&state_layout);
538
539        let state_clone = Rc::clone(&state);
540        let inner = move |tracers: &[Array]| -> Result<Vec<Array>, Exception> {
541            // put the tracers in their appropriate places:
542            // - arguments to the function
543            // - inner state
544
545            let tracer_args = &tracers[..args_len];
546
547            // save a snapshot of the inner state
548            let saved_state_inputs = state_snapshot(&mut **state_clone.borrow_mut())?;
549
550            // replace the inner state with the tracers
551            let state_input_count = saved_state_inputs.present_values().count();
552            if tracers.len() != args_len + state_input_count {
553                return Err(Exception::custom(format!(
554                    "compile_with_state: invalid tracer count - expected {} arguments and {} state inputs, got {} total inputs",
555                    args_len,
556                    state_input_count,
557                    tracers.len()
558                )));
559            }
560            let tracer_state = StateSnapshot::from_layout_and_values(
561                &saved_state_inputs.layout(),
562                &tracers[args_len..],
563            )?;
564            state_clone
565                .borrow_mut()
566                .state_projection()?
567                .restore(tracer_state, false)?;
568
569            // call the function with the tracer arguments and the state holding tracers
570            let mut result = (f)(*state_clone.borrow_mut(), tracer_args);
571
572            // Capture function output count before appending state
573            num_fn_outputs_clone.set(Some(result.len()));
574
575            // recapture the state as it may have changed
576            let state_output = state_snapshot(&mut **state_clone.borrow_mut())?;
577            let mut state_output_tracers =
578                state_output.present_values().cloned().collect::<Vec<_>>();
579
580            if state_layout_clone.borrow().is_none() {
581                *state_layout_clone.borrow_mut() = Some(state_output.layout());
582            }
583
584            // put the original values back in the state
585            restore_state(*state_clone.borrow_mut(), &saved_state_inputs, true)?;
586
587            // return the result of the function and the state
588            result.append(&mut state_output_tracers);
589
590            Ok(result)
591        };
592
593        let inner_closure = Closure::new_fallible(inner);
594        let result = call_mut_with_state_inner(
595            inner_closure,
596            self.id,
597            self.shapeless,
598            Rc::clone(&state),
599            args,
600            Rc::clone(&num_function_outputs),
601            Rc::clone(&state_layout),
602        );
603        self.num_function_outputs = num_function_outputs.get();
604        self.state_layout = state_layout.borrow().clone();
605        if let Err(error) = &result {
606            let after_failure = state_snapshot(&mut **state.borrow_mut())?;
607            if allow_growth_recovery && was_untraced && state_grew(&saved_state, &after_failure) {
608                restore_state(*state.borrow_mut(), &saved_state, true)?;
609                let retry =
610                    self.call_mut_with_state_attempt(&mut **state.borrow_mut(), args, false);
611                if retry.is_err() {
612                    restore_state(*state.borrow_mut(), &saved_state, false)?;
613                    self.cache.erase(self.id);
614                    self.num_function_outputs = None;
615                    self.state_layout = None;
616                }
617                return retry;
618            }
619            restore_state(*state.borrow_mut(), &saved_state, false).map_err(|restore_error| {
620                Exception::custom(format!(
621                    "{}; transactional restore failed: {}",
622                    error.what(),
623                    restore_error.what()
624                ))
625            })?;
626        }
627        result
628    }
629
630    fn fallible_call_mut_with_state<U>(
631        &mut self,
632        state: &mut U,
633        args: &[impl AsRef<Array>],
634    ) -> Result<Vec<Array>, Exception>
635    where
636        F: FnMut(&mut U, &[Array]) -> Result<Vec<Array>, Exception>,
637        U: Updatable,
638    {
639        self.fallible_call_mut_with_state_attempt(state, args, true)
640    }
641
642    fn fallible_call_mut_with_state_attempt<U>(
643        &mut self,
644        state: &mut U,
645        args: &[impl AsRef<Array>],
646        allow_growth_recovery: bool,
647    ) -> Result<Vec<Array>, Exception>
648    where
649        F: FnMut(&mut U, &[Array]) -> Result<Vec<Array>, Exception>,
650        U: Updatable,
651    {
652        let was_untraced = self.state_layout.is_none();
653        if let Some(expected) = self.state_layout.as_deref() {
654            validate_state_layout(expected, state, "call input")?;
655        }
656        let args_len = args.len();
657        let saved_state = state_snapshot(state)?;
658        let state = Rc::new(RefCell::new(state));
659        let f = &mut self.f;
660
661        // Cell to capture the number of function outputs during tracing
662        let num_function_outputs = Rc::new(Cell::new(self.num_function_outputs));
663        let num_fn_outputs_clone = Rc::clone(&num_function_outputs);
664        let state_layout = Rc::new(RefCell::new(self.state_layout.clone()));
665        let state_layout_clone = Rc::clone(&state_layout);
666
667        let state_clone = Rc::clone(&state);
668        let inner = move |tracers: &[Array]| -> Result<Vec<Array>, Exception> {
669            // put the tracers in their appropriate places:
670            // - arguments to the function
671            // - inner state
672
673            let tracer_args = &tracers[..args_len];
674
675            // save a snapshot of the inner state
676            let saved_state_inputs = state_snapshot(&mut **state_clone.borrow_mut())?;
677
678            // replace the inner state with the tracers
679            let state_input_count = saved_state_inputs.present_values().count();
680            if tracers.len() != args_len + state_input_count {
681                return Err(Exception::custom(format!(
682                    "compile_with_state: invalid tracer count - expected {} arguments and {} state inputs, got {} total inputs",
683                    args_len,
684                    state_input_count,
685                    tracers.len()
686                )));
687            }
688            let tracer_state = StateSnapshot::from_layout_and_values(
689                &saved_state_inputs.layout(),
690                &tracers[args_len..],
691            )?;
692            state_clone
693                .borrow_mut()
694                .state_projection()?
695                .restore(tracer_state, false)?;
696
697            // call the function with the tracer arguments and the state holding tracers
698            let call_result = {
699                let mut state = state_clone.borrow_mut();
700                (f)(&mut **state, tracer_args)
701            };
702            let mut result = match call_result {
703                Ok(result) => result,
704                Err(error) => {
705                    restore_state(*state_clone.borrow_mut(), &saved_state_inputs, false).map_err(
706                        |restore_error| {
707                            Exception::custom(format!(
708                                "{}; transactional restore failed: {}",
709                                error.what(),
710                                restore_error.what()
711                            ))
712                        },
713                    )?;
714                    return Err(error);
715                }
716            };
717
718            // Capture function output count before appending state
719            num_fn_outputs_clone.set(Some(result.len()));
720
721            // recapture the state as it may have changed
722            let state_output = state_snapshot(&mut **state_clone.borrow_mut())?;
723            let mut state_output_tracers =
724                state_output.present_values().cloned().collect::<Vec<_>>();
725
726            if state_layout_clone.borrow().is_none() {
727                *state_layout_clone.borrow_mut() = Some(state_output.layout());
728            }
729
730            // put the original values back in the state
731            restore_state(*state_clone.borrow_mut(), &saved_state_inputs, true)?;
732
733            // return the result of the function and the state
734            result.append(&mut state_output_tracers);
735
736            Ok(result)
737        };
738
739        let inner_closure = Closure::new_fallible(inner);
740        let result = call_mut_with_state_inner(
741            inner_closure,
742            self.id,
743            self.shapeless,
744            Rc::clone(&state),
745            args,
746            Rc::clone(&num_function_outputs),
747            Rc::clone(&state_layout),
748        );
749        self.num_function_outputs = num_function_outputs.get();
750        self.state_layout = state_layout.borrow().clone();
751        if let Err(error) = &result {
752            let after_failure = state_snapshot(&mut **state.borrow_mut())?;
753            if allow_growth_recovery && was_untraced && state_grew(&saved_state, &after_failure) {
754                restore_state(*state.borrow_mut(), &saved_state, true)?;
755                let retry = self.fallible_call_mut_with_state_attempt(
756                    &mut **state.borrow_mut(),
757                    args,
758                    false,
759                );
760                if retry.is_err() {
761                    restore_state(*state.borrow_mut(), &saved_state, false)?;
762                    self.cache.erase(self.id);
763                    self.num_function_outputs = None;
764                    self.state_layout = None;
765                }
766                return retry;
767            }
768            restore_state(*state.borrow_mut(), &saved_state, false).map_err(|restore_error| {
769                Exception::custom(format!(
770                    "{}; transactional restore failed: {}",
771                    error.what(),
772                    restore_error.what()
773                ))
774            })?;
775        }
776        result
777    }
778}
mlx_rs/transforms/compile/compile_with_state.rs

mlx_rs/transforms/compile/
compile_with_state.rs
