rlx_models_core/
lm.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
16//! Shared causal-LM flow helpers — re-export tier-0 surface for model authors.
17//!
18//! Model authors should use [`rlx_flow::prelude`] and arch recipes (`llama32::Llama32Flow`).
19//! Import HIR/Graph only via [`into_compile_parts`] at the compile boundary.
20
21pub use rlx_flow::prelude::*;
22
23pub trait FlowBuildExt {
24    fn with_lm_head(self, vocab: usize, hidden: usize, tied: bool) -> Self;
25    fn logits(self) -> Self;
26}
27
28impl FlowBuildExt for ModelFlow {
29    fn with_lm_head(self, vocab: usize, hidden: usize, tied: bool) -> Self {
30        self.lm_head(vocab, hidden, tied)
31    }
32
33    fn logits(self) -> Self {
34        self.output("logits")
35    }
36}
37
38/// Split built flow for compile — no Graph/HIR imports needed at call site.
39pub fn into_compile_parts(
40    built: BuiltModel,
41) -> anyhow::Result<(
42    rlx_ir::hir::HirModule,
43    std::collections::HashMap<String, Vec<f32>>,
44)> {
45    built.into_parts()
46}
rlx_models_core/lm.rs

rlx_models_core/
lm.rs
