mlx_c_rust/compile/
multi_output_graph.rs

1use crate::MlxCError;
2use crate::compile::graph::{MlxCompiledGraph, MlxGraphBuilder};
3use crate::{MlxArray, MlxBindingsContext};
4
5const APPLY_COMPILED_MULTI_OUTPUT_GRAPH_OPERATION: &str = "apply a compiled multi-output MLX graph";
6
7/// Owns one compiled MLX graph that produces an ordered vector of outputs.
8///
9/// Model-side compiled composites that need several outputs at once — the
10/// multi-token-prediction verification window returns successor states beside
11/// its logits — construct one of these with their own graph builder and apply
12/// it through [`MlxBindingsContext::apply_compiled_multi_output_graph`]. Output
13/// order is the builder's write order and never changes after compilation.
14#[derive(Debug)]
15pub struct MlxCompiledMultiOutputGraph {
16    compiled_graph: MlxCompiledGraph,
17}
18
19impl MlxCompiledMultiOutputGraph {
20    /// Compiles the given graph builder into a replayable multi-output graph.
21    /// The `shapeless` flag carries the same shape-specialization tradeoff as
22    /// the single-output compiled graph constructor.
23    pub fn new(
24        graph_builder: MlxGraphBuilder,
25        compile_operation: &'static str,
26        shapeless: bool,
27    ) -> Result<Self, MlxCError> {
28        Ok(Self {
29            compiled_graph: MlxCompiledGraph::new(graph_builder, compile_operation, shapeless)?,
30        })
31    }
32}
33
34impl MlxBindingsContext {
35    /// Applies a compiled multi-output graph and returns its ordered outputs.
36    pub fn apply_compiled_multi_output_graph(
37        &self,
38        compiled_graph: &MlxCompiledMultiOutputGraph,
39        graph_inputs: &[&MlxArray],
40    ) -> Result<Vec<MlxArray>, MlxCError> {
41        compiled_graph
42            .compiled_graph
43            .apply_multi(graph_inputs, APPLY_COMPILED_MULTI_OUTPUT_GRAPH_OPERATION)
44    }
45}
mlx_c_rust/compile/multi_output_graph.rs

mlx_c_rust/compile/
multi_output_graph.rs
