rlx_flow/blocks/
embed_scale.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3//
4// This program is free software: you can redistribute it and/or modify
5// it under the terms of the GNU General Public License as published by
6// the Free Software Foundation, version 3.
7//
8// This program is distributed in the hope that it will be useful,
9// but WITHOUT ANY WARRANTY; without even the implied warranty of
10// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
11// GNU General Public License for more details.
12//
13// You should have received a copy of the GNU General Public License
14// along with this program. If not, see <https://www.gnu.org/licenses/>.
15
16use anyhow::Result;
17use rlx_ir::HirGraphExt;
18use rlx_ir::hir::HirMut;
19use rlx_ir::{DType, Shape};
20
21use super::BlockStage;
22use crate::context::FlowCtx;
23use crate::value::FlowValue;
24
25/// Multiply token embeddings by `sqrt(hidden_size)` (Gemma / Gemma 2).
26#[derive(Debug, Clone)]
27pub struct EmbedScaleStage {
28    pub hidden_size: usize,
29}
30
31impl EmbedScaleStage {
32    pub fn new(hidden_size: usize) -> Self {
33        Self { hidden_size }
34    }
35}
36
37impl BlockStage for EmbedScaleStage {
38    fn emit(&self, ctx: &mut FlowCtx<'_>, input: FlowValue) -> Result<Option<FlowValue>> {
39        let scale = (self.hidden_size as f32).sqrt();
40        let name = format!("gemma.embed_scale.{}", self.hidden_size);
41        let scale_id = ctx.synth_param(&name, vec![scale], Shape::new(&[1], DType::F32));
42        let mut gb = HirMut::new(ctx.hir());
43        let out = gb.mul(input.id, scale_id);
44        Ok(Some(ctx.wrap(out, input.shape.clone())))
45    }
46}
rlx_flow/blocks/embed_scale.rs

rlx_flow/blocks/
embed_scale.rs
