Skip to main content

compile

Function compile 

Source
pub fn compile<'f, I, O, F, R>(f: F) -> impl Fn(&I) -> Result<O> + 'f
where I: Tree, O: Tree, F: Fn(&I) -> R + 'f, R: IntoResult<Ok = O>,
Expand description

Returns f compiled: a function with the same results that traces f once per input structure (the TreeDef: paths, leaf kinds, trainable bits, enum variants, mode flags and static fields), leaf dtypes and shapes, and default stream, and replays the simplified, fused graph on later calls.

f takes a tree (an array, a tuple, a struct) and returns one, directly or as a Result; it may borrow locals (the cache stores tapes, not f). Write the closure’s parameter type (|x: &Array|).

As in MLX, f must be pure:

  • arrays it captures are constants, frozen at the first trace (so is random state drawn inside it; compile_mut threads the random key instead);
  • arrays cannot be evaluated inside it.

Called on arrays being traced by a transform (grad, vmap, another compile), or on failed arrays, it runs f directly, so transforms see the original operations and failures stay sticky.

Each compiled function keeps the tapes of its 64 most recently used signatures (Options::cache_capacity). A trace whose constants include an array being traced by a running transform (a tracer f captures) is used once and not kept.