rlx_rl/graph/
mlp.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// RLX — shared MLP helpers.
16
17use rlx_ir::infer::GraphExt;
18use rlx_ir::{DType, Graph, NodeId, Shape};
19
20/// Trainable parameter (forward + backward grad node).
21#[derive(Debug, Clone)]
22pub struct ParamSlot {
23    pub name: String,
24    pub shape: Vec<usize>,
25    pub param: NodeId,
26    pub grad: Option<NodeId>,
27}
28
29impl ParamSlot {
30    pub fn numel(&self) -> usize {
31        self.shape.iter().product()
32    }
33}
34
35/// Hidden MLP trunk (no output head). Returns `(hidden, out_dim)`.
36pub fn mlp_trunk(
37    g: &mut Graph,
38    x: NodeId,
39    in_dim: usize,
40    hidden: &[usize],
41    prefix: &str,
42    params: &mut Vec<ParamSlot>,
43) -> (NodeId, usize) {
44    let f = DType::F32;
45    let mut h = x;
46    let mut in_d = in_dim;
47
48    for (li, &hd) in hidden.iter().enumerate() {
49        let w_name = format!("{prefix}_w{li}");
50        let b_name = format!("{prefix}_b{li}");
51        let w = g.param(&w_name, Shape::new(&[in_d, hd], f));
52        let b = g.param(&b_name, Shape::new(&[hd], f));
53        params.push(ParamSlot {
54            name: w_name,
55            shape: vec![in_d, hd],
56            param: w,
57            grad: None,
58        });
59        params.push(ParamSlot {
60            name: b_name,
61            shape: vec![hd],
62            param: b,
63            grad: None,
64        });
65        h = g.mm(h, w);
66        h = g.add(h, b);
67        h = g.tanh(h);
68        in_d = hd;
69    }
70    (h, in_d)
71}
72
73/// Full MLP including linear head to `out_dim`.
74pub fn mlp_layers(
75    g: &mut Graph,
76    x: NodeId,
77    in_dim: usize,
78    hidden: &[usize],
79    out_dim: usize,
80    prefix: &str,
81    params: &mut Vec<ParamSlot>,
82) -> NodeId {
83    let f = DType::F32;
84    let (h, in_d) = mlp_trunk(g, x, in_dim, hidden, prefix, params);
85    let w_name = format!("{prefix}_w_out");
86    let b_name = format!("{prefix}_b_out");
87    let w = g.param(&w_name, Shape::new(&[in_d, out_dim], f));
88    let b = g.param(&b_name, Shape::new(&[out_dim], f));
89    params.push(ParamSlot {
90        name: w_name,
91        shape: vec![in_d, out_dim],
92        param: w,
93        grad: None,
94    });
95    params.push(ParamSlot {
96        name: b_name,
97        shape: vec![out_dim],
98        param: b,
99        grad: None,
100    });
101    let out = g.mm(h, w);
102    g.add(out, b)
103}
104
105pub fn flow_map_jump(
106    g: &mut Graph,
107    a_r: NodeId,
108    u: NodeId,
109    r: NodeId,
110    t: NodeId,
111    batch: usize,
112) -> NodeId {
113    let dt = g.sub(t, r);
114    let dt_col = g.reshape_(dt, vec![batch as i64, 1]);
115    let scaled = g.mul(u, dt_col);
116    g.add(a_r, scaled)
117}
118
119pub fn mse_mean(g: &mut Graph, pred: NodeId, target: NodeId) -> NodeId {
120    let diff = g.sub(pred, target);
121    let sq = g.mul(diff, diff);
122    let rank = g.shape(sq).rank();
123    let axes: Vec<usize> = (0..rank).collect();
124    g.mean(sq, axes, false)
125}
126
127pub fn concat_features(g: &mut Graph, parts: Vec<NodeId>) -> NodeId {
128    g.concat_(parts, 1)
129}
130
131pub fn init_mat(
132    w: &mut crate::graph::actor::WeightStore,
133    name: &str,
134    rows: usize,
135    cols: usize,
136    seed: &mut u64,
137) {
138    let scale = (2.0 / (rows + cols) as f32).sqrt();
139    let n = rows * cols;
140    let mut v = vec![0.0f32; n];
141    for x in &mut v {
142        *seed = crate::buffer::rand_like(*seed);
143        let u = (*seed >> 11) as f32 / (1u32 << 21) as f32;
144        *seed = crate::buffer::rand_like(*seed);
145        let n2 = (*seed >> 11) as f32 / (1u32 << 21) as f32;
146        *x = (u * 2.0 * std::f32::consts::PI * n2).sin() * scale;
147    }
148    w.0.insert(name.to_string(), v);
149}
150
151pub fn init_vec(w: &mut crate::graph::actor::WeightStore, name: &str, n: usize, seed: &mut u64) {
152    let mut v = vec![0.0f32; n];
153    for x in &mut v {
154        *seed = crate::buffer::rand_like(*seed);
155        *x = 0.01 * ((*seed >> 11) as f32 / (1u32 << 21) as f32 - 0.5);
156    }
157    w.0.insert(name.to_string(), v);
158}
rlx_rl/graph/mlp.rs

rlx_rl/graph/
mlp.rs
