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, ) -> ModelFlow

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

pub fn profile_encoder(self) -> ModelFlow

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>, ) -> ModelFlow

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>, ) -> ModelFlow

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, ) -> ModelFlow

LayerNorm with separate gamma/beta weights.

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

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, ) -> ModelFlow

Repeat NomicBERT encoder layers.

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

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, ) -> ModelFlow

Repeat BERT encoder layers with auto-named prefixes.

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

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, ) -> ModelFlow

Repeat DINOv2 ViT encoder blocks.

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

Repeat NomicVision encoder blocks.

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

Compatibility shim: repeat SigLIP-style vision layers.

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

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

pub fn profile_prefill(self) -> ModelFlow

Fusion-first prefill defaults.

pub fn profile_decode(self) -> ModelFlow

Decode / KV-cache defaults (Fusable lowering).

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

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

pub fn token_embed(self) -> ModelFlow

HuggingFace-style token embedding table.

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

Precomputed RoPE sin/cos tables stored as params.

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

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

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

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

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, ) -> ModelFlow

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

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

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

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

Build a named layer from a LayerStack closure.

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

Fused LLaMA prefill layer (default fast path).

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

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

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

pub fn residual_save(self) -> ModelFlow

pub fn residual_add(self) -> ModelFlow

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

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

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

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

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

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

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

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, ) -> ModelFlow

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

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

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

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

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

pub fn hidden_states(self) -> ModelFlow

Hidden states output (no LM head).

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

LLaMA prefill decoder block at layer_idx.

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

LLaMA decode block with KV-cache concat.

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

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

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

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

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

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

Gather last token (dynamic last_token_idx input).

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

Gather last token at fixed sequence length.

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

Causal LM head — tied or separate weights.

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

Tier-2 escape hatch — append a raw stage.

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

Append multiple raw stages in order.

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

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

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

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

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

where F: Fn(&mut Emit<'_>, Option<FlowValue>) -> Result<Option<FlowValue>, Error> + 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) -> ModelFlow

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

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

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

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

impl ModelFlow

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

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

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

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

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

pub fn profile(&self) -> &CompileProfile

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

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

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

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

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

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

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

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

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

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<(), Error>

Formats the value using the given formatter.

impl FlowBuildExt for ModelFlow

fn with_lm_head(self, vocab: usize, hidden: usize, tied: bool) -> Self

fn logits(self) -> Self