pub fn compile<'f, I, O, F, R>(f: F) -> impl Fn(&I) -> Result<O> + 'fExpand 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_mutthreads 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.