rlx-ir 0.2.0

Tensor IR for the RLX ML compiler — standalone, serializable, optimizable
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114

// RLX — versatile ML compiler + runtime.
// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, version 3.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program. If not, see <https://www.gnu.org/licenses/>.

//! Normalization builders: LayerNorm, fused residual+LN (plan #53).

use crate::{Graph, NodeId, Op, Shape};

impl Graph {
    /// LayerNorm2d on NCHW (normalize across channels at each spatial position).
    pub fn layer_norm2d(&mut self, input: NodeId, gamma: NodeId, beta: NodeId, eps: f32) -> NodeId {
        let shape = self.node(input).shape.clone();
        self.push(
            Op::LayerNorm2d { eps },
            vec![input, gamma, beta],
            shape,
            None,
        )
    }

    /// Group normalization on NCHW.
    pub fn group_norm(
        &mut self,
        input: NodeId,
        gamma: NodeId,
        beta: NodeId,
        num_groups: usize,
        eps: f32,
    ) -> NodeId {
        let shape = self.node(input).shape.clone();
        self.push(
            Op::GroupNorm { num_groups, eps },
            vec![input, gamma, beta],
            shape,
            None,
        )
    }

    /// Layer normalization.
    pub fn layer_norm(
        &mut self,
        input: NodeId,
        gamma: NodeId,
        beta: NodeId,
        axis: i32,
        eps: f32,
        shape: Shape,
    ) -> NodeId {
        self.push(
            Op::LayerNorm { axis, eps },
            vec![input, gamma, beta],
            shape,
            None,
        )
    }

    /// Fused residual + bias + layer norm (created by optimization passes).
    pub fn fused_residual_ln(
        &mut self,
        x: NodeId,
        residual: NodeId,
        bias: Option<NodeId>,
        gamma: NodeId,
        beta: NodeId,
        eps: f32,
        shape: Shape,
    ) -> NodeId {
        let has_bias = bias.is_some();
        let mut inputs = vec![x, residual];
        if let Some(b) = bias {
            inputs.push(b);
        }
        inputs.push(gamma);
        inputs.push(beta);
        self.push(Op::FusedResidualLN { has_bias, eps }, inputs, shape, None)
    }

    /// Fused residual + bias + RMS norm (created by optimization passes).
    pub fn fused_residual_rms_norm(
        &mut self,
        x: NodeId,
        residual: NodeId,
        bias: Option<NodeId>,
        gamma: NodeId,
        beta: NodeId,
        eps: f32,
        shape: Shape,
    ) -> NodeId {
        let has_bias = bias.is_some();
        let mut inputs = vec![x, residual];
        if let Some(b) = bias {
            inputs.push(b);
        }
        inputs.push(gamma);
        inputs.push(beta);
        self.push(
            Op::FusedResidualRmsNorm { has_bias, eps },
            inputs,
            shape,
            None,
        )
    }
}