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, compiled as MLX compiles its: scalar constants are literals in the source, and Metal’s default (fast) math functions apply, which only power uses (the other functions are the precise ones, as in MLX). Fused results equal the uncompiled function’s bit for bit, except where power is fused: there they can differ in the low bits, as mx.compile’s do. 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

With Options::shapeless a function is traced once per input structure, dtypes and ranks, and replayed for any shapes (MLX’s shapeless=True).

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, compile_mut and compile_flat_with: 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_flat_with

compile_flat with options (Options::shapeless, the cache capacity).

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)).