rlx_flow/blocks/
embed_scale.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3
4use anyhow::Result;
5use rlx_ir::HirGraphExt;
6use rlx_ir::hir::HirMut;
7use rlx_ir::{DType, Shape};
8
9use super::BlockStage;
10use crate::context::FlowCtx;
11use crate::value::FlowValue;
12
13/// Multiply token embeddings by `sqrt(hidden_size)` (Gemma / Gemma 2).
14#[derive(Debug, Clone)]
15pub struct EmbedScaleStage {
16    pub hidden_size: usize,
17}
18
19impl EmbedScaleStage {
20    pub fn new(hidden_size: usize) -> Self {
21        Self { hidden_size }
22    }
23}
24
25impl BlockStage for EmbedScaleStage {
26    fn emit(&self, ctx: &mut FlowCtx<'_>, input: FlowValue) -> Result<Option<FlowValue>> {
27        let scale = (self.hidden_size as f32).sqrt();
28        let name = format!("gemma.embed_scale.{}", self.hidden_size);
29        let scale_id = ctx.synth_param(&name, vec![scale], Shape::new(&[1], DType::F32));
30        let mut gb = HirMut::new(ctx.hir());
31        let out = gb.mul(input.id, scale_id);
32        Ok(Some(ctx.wrap(out, input.shape.clone())))
33    }
34}
rlx_flow/blocks/embed_scale.rs

rlx_flow/blocks/
embed_scale.rs
