Skip to main content

compile

Function compile 

Source
pub fn compile<F>(f: F) -> impl Fn(&[Array]) -> Result<Vec<Array>>
where F: Fn(&[Array]) -> Result<Vec<Array>>,
Expand description

Returns f compiled: a function with the same results that traces f once per input signature (dtypes and shapes) and default stream, and replays the simplified, fused graph on later calls.

As in MLX, f must be pure:

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

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