rlx_flow/blocks/
cls_pool.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::hir::HirMut;
18use rlx_ir::{DType, HirGraphExt};
19
20use super::BlockStage;
21use crate::context::FlowCtx;
22use crate::value::FlowValue;
23
24/// Extract CLS token `[batch, 1, hidden]` → `[batch, hidden]`.
25#[derive(Debug, Clone)]
26pub struct ClsTokenPoolStage {
27    pub batch: usize,
28    pub hidden: usize,
29}
30
31impl ClsTokenPoolStage {
32    pub fn new(batch: usize, hidden: usize) -> Self {
33        Self { batch, hidden }
34    }
35}
36
37impl BlockStage for ClsTokenPoolStage {
38    fn emit(&self, ctx: &mut FlowCtx<'_>, input: FlowValue) -> Result<Option<FlowValue>> {
39        let mut gb = HirMut::new(ctx.hir());
40        let cls = gb.narrow_(input.id, 1, 0, 1);
41        let flat = gb.reshape_(cls, vec![self.batch as i64, self.hidden as i64]);
42        Ok(Some(ctx.wrap(
43            flat,
44            rlx_ir::Shape::new(&[self.batch, self.hidden], DType::F32),
45        )))
46    }
47}
rlx_flow/blocks/cls_pool.rs

rlx_flow/blocks/
cls_pool.rs
