Struct GemmaFlow 

pub struct GemmaFlow<'a> { /* private fields */ }

Fluent Gemma flow builder — reads config once, chain modifiers, then build.

use rlx_models::gemma::{GemmaConfig, GemmaFlow};

let built = GemmaFlow::new(&cfg)
    .prefill()
    .batch(1)
    .seq(128)
    .lm_head()
    .last_token_logits()
    .build(&mut weights)?;

Implementations

impl<'a> GemmaFlow<'a>

pub fn new(cfg: &'a GemmaConfig) -> GemmaFlow<'a>

pub fn prefill_from_hidden(self) -> GemmaFlow<'a>

Skip token embedding — feed pre-scaled hidden states at prefill_hidden.

pub fn prefill_media_attn_bias(self) -> GemmaFlow<'a>

Add attn_bias input and bidirectional self-attn on sliding layers.

pub fn for_prefill( cfg: &'a GemmaConfig, batch: usize, seq: usize, ) -> GemmaFlow<'a>

Prefill recipe with common batch/seq defaults.

pub fn for_decode( cfg: &'a GemmaConfig, batch: usize, past_seq: usize, ) -> GemmaFlow<'a>

Decode recipe with common batch/past defaults (includes LM head).

pub fn prefill(self) -> GemmaFlow<'a>

pub fn decode(self) -> GemmaFlow<'a>

pub fn batch(self, batch: usize) -> GemmaFlow<'a>

pub fn seq(self, seq: usize) -> GemmaFlow<'a>

Prefill sequence length (ignored in decode mode).

pub fn past(self, past_seq: usize) -> GemmaFlow<'a>

Decode past length (ignored in prefill mode).

pub fn dynamic_seq(self) -> GemmaFlow<'a>

Symbolic sequence dim (sym::SEQ) for dynamic prefill specialization.

pub fn dynamic_past(self) -> GemmaFlow<'a>

Symbolic past dim (sym::PAST_SEQ) for dynamic decode specialization.

pub fn lm_head(self) -> GemmaFlow<'a>

pub fn hidden_only(self) -> GemmaFlow<'a>

Hidden states only — skip LM head (default for prefill unless .lm_head()).

pub fn last_token_logits(self) -> GemmaFlow<'a>

pub fn export_kv(self) -> GemmaFlow<'a>

pub fn custom_mask(self) -> GemmaFlow<'a>

pub fn profile(self, profile: CompileProfile) -> GemmaFlow<'a>

pub fn profile_prefill(self) -> GemmaFlow<'a>

Fusion-first prefill profile preset.

pub fn profile_decode(self) -> GemmaFlow<'a>

pub fn profile_near(self, weights_path: &Path) -> GemmaFlow<'a>

pub fn before_layers( self, stages: impl IntoIterator<Item = FlowStage>, ) -> GemmaFlow<'a>

Insert custom stages after embedding, before the layer stack.

pub fn after_layers( self, stages: impl IntoIterator<Item = FlowStage>, ) -> GemmaFlow<'a>

Insert custom stages after the layer stack, before final norm / LM head.

pub fn layer<F>(self, f: F) -> GemmaFlow<'a>

where F: Fn(GemmaLayerCtx<'_>) -> FlowStage + Send + Sync + 'static,

Override per-layer construction (prefill or decode depending on mode).

Call GemmaLayerCtx::default_stage to keep stock blocks for unmodified layers.

pub fn patch_flow<F>(self, f: F) -> GemmaFlow<'a>

where F: Fn(ModelFlow) -> ModelFlow + Send + Sync + 'static,

Patch the assembled ModelFlow before build — full flexibility escape hatch.

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

Trait Implementations

impl<'a> Clone for GemmaFlow<'a>

fn clone(&self) -> GemmaFlow<'a>

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Debug for GemmaFlow<'_>

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

Formats the value using the given formatter.