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

rlx_flow/blocks/
embed.rs
