rlx_flow/blocks/
cls_pool.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3
4use anyhow::Result;
5use rlx_ir::hir::HirMut;
6use rlx_ir::{DType, HirGraphExt};
7
8use super::BlockStage;
9use crate::context::FlowCtx;
10use crate::value::FlowValue;
11
12/// Extract CLS token `[batch, 1, hidden]` → `[batch, hidden]`.
13#[derive(Debug, Clone)]
14pub struct ClsTokenPoolStage {
15    pub batch: usize,
16    pub hidden: usize,
17}
18
19impl ClsTokenPoolStage {
20    pub fn new(batch: usize, hidden: usize) -> Self {
21        Self { batch, hidden }
22    }
23}
24
25impl BlockStage for ClsTokenPoolStage {
26    fn emit(&self, ctx: &mut FlowCtx<'_>, input: FlowValue) -> Result<Option<FlowValue>> {
27        let mut gb = HirMut::new(ctx.hir());
28        let cls = gb.narrow_(input.id, 1, 0, 1);
29        let flat = gb.reshape_(cls, vec![self.batch as i64, self.hidden as i64]);
30        Ok(Some(ctx.wrap(
31            flat,
32            rlx_ir::Shape::new(&[self.batch, self.hidden], DType::F32),
33        )))
34    }
35}
rlx_flow/blocks/cls_pool.rs

rlx_flow/blocks/
cls_pool.rs
