rlx_flow/blocks/
embed.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;
19
20use super::BlockStage;
21use crate::context::FlowCtx;
22use crate::value::FlowValue;
23#[derive(Debug, Clone)]
24pub struct EmbedStage {
25    pub weight_key: String,
26    pub axis: usize,
27}
28
29impl EmbedStage {
30    pub fn token(weight_key: impl Into<String>) -> Self {
31        Self {
32            weight_key: weight_key.into(),
33            axis: 0,
34        }
35    }
36}
37
38impl BlockStage for EmbedStage {
39    fn emit(&self, ctx: &mut FlowCtx<'_>, input: FlowValue) -> Result<Option<FlowValue>> {
40        let embed_w = ctx.load_param(&self.weight_key, false)?;
41        ctx.state.embed_weight = Some(embed_w);
42        let out_shape = {
43            let w_shape = ctx.hir().node(embed_w).shape.clone();
44            let mut dims: Vec<rlx_ir::Dim> = input.shape.dims().to_vec();
45            dims.push(w_shape.dim(1));
46            rlx_ir::Shape::from_dims(&dims, input.shape.dtype())
47        };
48        let mut gb = HirMut::new(ctx.hir());
49        let id = gb.gather_(embed_w, input.id, self.axis);
50        Ok(Some(ctx.wrap(id, out_shape)))
51    }
52}
rlx_flow/blocks/embed.rs

rlx_flow/blocks/
embed.rs
