Struct ModelFlow 

pub struct ModelFlow { /* private fields */ }

Block assembly-line builder — tier-0 model author surface.

Implementations

impl ModelFlow

pub fn profile_file( self, path: impl AsRef<Path>, default: fn() -> CompileProfile, ) -> Self

Load tier-1 profile from a *.rlx.toml file (falls back to default on error).

pub fn profile_encoder(self) -> Self

Encoder / embedding model defaults (Direct lowering, no KV fusion).

pub fn gather_from_input( self, input_name: impl Into<String>, weight_key: impl Into<String>, ) -> Self

Gather rows from a side input into the primary flow (starts embedding stack).

pub fn gather_add( self, input_name: impl Into<String>, weight_key: impl Into<String>, ) -> Self

Add an embedding looked up from a side input.

pub fn layer_norm( self, gamma_key: impl Into<String>, beta_key: impl Into<String>, eps: f32, ) -> Self

LayerNorm with separate gamma/beta weights.

pub fn gelu_ffn(self, layer_prefix: impl Into<String>) -> Self

BERT-style GELU FFN under a layer prefix.

pub fn repeat_nomic_layers( self, count: usize, hidden_size: usize, num_heads: usize, head_dim: usize, eps: f32, ) -> Self

Repeat NomicBERT encoder layers.

pub fn bert_encoder_layer(self, spec: BertEncoderLayerSpec) -> Self

BERT-style encoder layer (fused QKV + padding-mask attention + GELU FFN).

pub fn repeat_bert_layers( self, count: usize, prefix: impl Into<String>, qkv_style: BertQkvStyle, hidden_size: usize, num_heads: usize, eps: f32, ) -> Self

Repeat BERT encoder layers with auto-named prefixes.

pub fn attn_mask_ones(self, batch: usize, seq: usize) -> Self

Synthesize an all-ones attention mask for vision encoders (no padding).

pub fn repeat_dinov2_layers( self, count: usize, hidden_size: usize, num_heads: usize, eps: f32, ) -> Self

Repeat DINOv2 ViT encoder blocks.

pub fn repeat_vision_layers( self, count: usize, hidden_size: usize, num_heads: usize, eps: f32, ) -> Self

Repeat NomicVision encoder blocks.

pub fn repeat_siglip_layers( self, count: usize, hidden_size: usize, num_heads: usize, eps: f32, ) -> Self

Compatibility shim: repeat SigLIP-style vision layers.

pub fn cls_token_pool(self, batch: usize, hidden: usize) -> Self

Pool CLS token: [batch, seq, hidden] → [batch, hidden].

pub fn profile_prefill(self) -> Self

Fusion-first prefill defaults.

pub fn profile_decode(self) -> Self

Decode / KV-cache defaults (Fusable lowering).

pub fn embed(self, weight_key: impl Into<String>) -> Self

Token embedding (model.embed_tokens.weight by default).

pub fn token_embed(self) -> Self

HuggingFace-style token embedding table.

pub fn rope_tables(self, tables: RopeTablesStage) -> Self

Precomputed RoPE sin/cos tables stored as params.

pub fn zero_beta(self, len: usize) -> Self

Rank-1 zero vector for RMSNorm beta slots (LLaMA has no beta).

pub fn zero_beta_named(self, name: impl Into<String>, len: usize) -> Self

pub fn bind_decode_inputs(self, num_layers: usize, custom_mask: bool) -> Self

Bind decode inputs (call after declaring rope_cos, past_k_*, …).

pub fn repeat_layers( self, count: usize, stage_for_layer: impl Fn(usize) -> FlowStage + Send + Sync + 'static, ) -> Self

Repeat a per-layer stage count times (layer index passed to closure).

pub fn named_layer(self, name: impl Into<String>, inner: FlowStage) -> Self

Named decoder layer (shows up in fusion / inspect dumps).

pub fn layer( self, name: impl Into<String>, build: impl FnOnce(LayerStack) -> LayerStack, ) -> Self

Build a named layer from a LayerStack closure.

pub fn llama_prefill_layer( self, layer_idx: usize, spec: LlamaDecoderSpec, ) -> Self

Fused LLaMA prefill layer (default fast path).

pub fn llama_prefill_layer_composed( self, layer_idx: usize, spec: LlamaDecoderSpec, ) -> Self

Composed LLaMA prefill layer (small blocks — customize via LayerStack).

pub fn linear(self, weight_key: impl Into<String>, transpose: bool) -> Self

pub fn residual_save(self) -> Self

pub fn residual_add(self) -> Self

pub fn swiglu( self, gate_key: impl Into<String>, up_key: impl Into<String>, down_key: impl Into<String>, ) -> Self

pub fn swiglu_hf_mlp(self, prefix: impl Into<String>) -> Self

pub fn self_attn_prefill(self, spec: SelfAttnPrefillSpec) -> Self

pub fn gdn_scan(self, stage: GdnScanStage) -> Self

pub fn store_stream(self, name: impl Into<String>) -> Self

pub fn load_stream(self, name: impl Into<String>) -> Self

pub fn bind_inputs_to_streams( self, pairs: impl IntoIterator<Item = (impl Into<String>, impl Into<String>)>, ) -> Self

Bind declared graph inputs into named streams (multi-input models).

Example: FLUX .bind_inputs_to_streams(&[("hidden", "img"), ("encoder", "txt")]).

pub fn dual_stream<F>( self, name: impl Into<String>, stream_a: impl Into<String>, stream_b: impl Into<String>, f: F, ) -> Self

where F: Fn(&mut Emit<'_>, FlowValue, FlowValue) -> Result<(FlowValue, FlowValue)> + Send + Sync + 'static,

pub fn plugin<F>(self, f: F) -> Self

where F: Fn(&mut Emit<'_>, Option<FlowValue>) -> Result<Option<FlowValue>> + Send + Sync + 'static,

pub fn plugin_named<F>(self, name: impl Into<String>, f: F) -> Self

where F: Fn(&mut Emit<'_>, Option<FlowValue>) -> Result<Option<FlowValue>> + Send + Sync + 'static,

pub fn hidden_states(self) -> Self

Hidden states output (no LM head).

pub fn llama_decoder_layer( self, layer_idx: usize, spec: LlamaDecoderSpec, ) -> Self

LLaMA prefill decoder block at layer_idx.

pub fn llama_decode_layer( self, layer_idx: usize, spec: LlamaDecodeLayerSpec, kv_out: SideOutputs, ) -> Self

LLaMA decode block with KV-cache concat.

pub fn llama_kv_tap( self, layer_idx: usize, head_dim: usize, eps: f32, sink: &SideOutputs, ) -> Self

Side-effect K/V tap before a prefill layer (exports cache tensors).

pub fn final_norm(self, eps: f32) -> Self

Final RMSNorm before LM head (model.norm.weight by default).

pub fn rms_norm(self, weight_key: impl Into<String>, eps: f32) -> Self

pub fn gather_last_token_dynamic(self, batch: usize) -> Self

Gather last token (dynamic last_token_idx input).

pub fn gather_last_token_at(self, batch: usize, seq: usize) -> Self

Gather last token at fixed sequence length.

pub fn lm_head( self, vocab_size: usize, hidden_size: usize, tie_word_embeddings: bool, ) -> Self

Causal LM head — tied or separate weights.

pub fn raw_stage(self, stage: FlowStage) -> Self

Tier-2 escape hatch — append a raw stage.

pub fn raw_stages(self, stages: impl IntoIterator<Item = FlowStage>) -> Self

Append multiple raw stages in order.

pub fn sequence(self, stages: impl IntoIterator<Item = FlowStage>) -> Self

Run a list of stages as one nested sequence (side-effect stages allowed).

pub fn when(self, cond: bool, f: impl FnOnce(Self) -> Self) -> Self

Conditionally transform the builder (e.g. optional vision tower).

pub fn custom<F>(self, f: F) -> Self

where F: Fn(&mut Emit<'_>, Option<FlowValue>) -> Result<Option<FlowValue>> + Send + Sync + 'static,

Tier-2 custom subgraph — prefer promoting repeated patterns to blocks.

pub fn custom_named<F>(self, name: impl Into<String>, f: F) -> Self

where F: Fn(&mut Emit<'_>, Option<FlowValue>) -> Result<Option<FlowValue>> + Send + Sync + 'static,

Named custom subgraph (shows up in fusion / inspect dumps).

pub fn patch(self, f: impl FnOnce(Self) -> Self) -> Self

Patch the builder after preset assembly (arch recipes, Llama32Flow hooks).

impl ModelFlow

pub fn new(name: impl Into<String>) -> Self

pub fn with_extensions(self, plan: FlowExtensionPlan) -> Self

HIR extensions to apply after assemble, before compile (retroactive plugins).

pub fn input(self, name: impl Into<String>, shape: Shape) -> Self

Declare a graph input. The first input starts the tensor flow; later inputs are side declarations only (e.g. last_token_idx).

pub fn with_profile(self, profile: CompileProfile) -> Self

pub fn profile(&self) -> &CompileProfile

pub fn stage(self, stage: FlowStage) -> Self

pub fn output(self, name: impl Into<String>) -> Self

pub fn outputs(self, names: impl IntoIterator<Item = impl Into<String>>) -> Self

pub fn with_extra_outputs(self, ids: Vec<HirNodeId>) -> Self

Append side outputs (e.g. per-layer KV taps) after the primary output.

pub fn from_recipe(recipe: &impl ModelRecipe) -> Self

Build from a composable recipe, then optionally patch before compile.

pub fn build(self, weights: &mut dyn WeightSource) -> Result<BuiltModel>

pub fn build_with( self, weights: &mut dyn WeightSource, _gguf_packed: Option<&GgufPackedParams>, ) -> Result<BuiltModel>

Compatibility shim: older callers passed GGUF packed matmul params.

The current flow builder ignores packed params; packed lowering lives in model crates.

Trait Implementations

impl Debug for ModelFlow

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.