Crate nuvec_compile 

Function compilation: tracing, graph simplification and kernel fusion.

compile works as MLX’s mx.compile. The first call with a given input structure and signature traces the function on placeholders and simplifies the graph:

• identity views and same-type conversions are removed (explicit copies and dependencies, Primitive::Copy and Primitive::Depends, are kept);
• equal scalar constants are merged;
• equal operations on equal inputs are computed once (common-subexpression elimination).

On the GPU, chains of element-wise operations are then fused into single generated kernels. Later calls replay the simplified graph on the new inputs without running the function again.

use nuvec_compile::compile;
use nuvec_core::{Array, Result, array};

// x * x + 1, as one kernel on the GPU
let f = compile(|x: &Array| x.multiply(x).add(1.0));
assert_eq!(f(&array![1.0, 2.0])?.to_vec::<f32>()?, [2.0, 5.0]);
assert_eq!(f(&array![3.0, 4.0])?.to_vec::<f32>()?, [10.0, 17.0]); // replayed

compile_mut compiles a stateful step (MLX’s mx.compile(step, inputs=state, outputs=state)): a training step over (model, optimizer) whose leaves are inputs and outputs of the tape.

Stability: no promise when this crate is depended on directly; depend on nuvec, whose re-exports are semver-tracked.

Status: pre-alpha.

Structs

Options

Options of compile and compile_mut: Options::default().cache_capacity(8).

Constants

VERSION

The version of this crate.

Functions

compile

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.

compile_flat

compile over positional arrays: traces once per input signature (dtypes and shapes) and default stream.

compile_mut

Compiles a stateful step, as MLX’s mx.compile(step, inputs=state, outputs=state): f(&mut state, &inputs) where state is a tree (typically (model, optimizer)).