rlx_flow/blocks/
layer_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;
7
8use super::BlockStage;
9use crate::context::FlowCtx;
10use crate::value::FlowValue;
11
12#[derive(Debug, Clone)]
13pub struct LayerScaleStage {
14    pub gamma_key: String,
15}
16
17impl LayerScaleStage {
18    pub fn new(gamma_key: impl Into<String>) -> Self {
19        Self {
20            gamma_key: gamma_key.into(),
21        }
22    }
23}
24
25impl BlockStage for LayerScaleStage {
26    fn emit(&self, ctx: &mut FlowCtx<'_>, input: FlowValue) -> Result<Option<FlowValue>> {
27        let gamma = ctx.load_param(&self.gamma_key, false)?;
28        let mut gb = HirMut::new(ctx.hir());
29        let out = gb.mul(input.id, gamma);
30        Ok(Some(ctx.wrap(out, input.shape.clone())))
31    }
32}
rlx_flow/blocks/layer_scale.rs

rlx_flow/blocks/
layer_scale.rs
