rlx-flow 0.2.1

Block assembly-line API for RLX model builders — fusion-first, config-driven
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27

// RLX — versatile ML compiler + runtime.
// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.

use anyhow::Result;
use rlx_ir::{DType, Shape};

use crate::context::FlowCtx;

/// Bake single-step decode RoPE cos/sin tables into flow state (static past length).
#[derive(Debug, Clone)]
pub struct DecodeRopeParamsStage {
    pub cos: Vec<f32>,
    pub sin: Vec<f32>,
    pub half_dim: usize,
}

impl DecodeRopeParamsStage {
    pub fn emit(&self, ctx: &mut FlowCtx<'_>) -> Result<()> {
        let f = DType::F32;
        let shape = Shape::new(&[1, self.half_dim], f);
        let cos_id = ctx.synth_param("decode.rope.cos", self.cos.clone(), shape.clone());
        let sin_id = ctx.synth_param("decode.rope.sin", self.sin.clone(), shape);
        ctx.state.rope_cos = Some(cos_id);
        ctx.state.rope_sin = Some(sin_id);
        Ok(())
    }
}