rlx_ir/ops/
shape_ops.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//! Shape-manipulation builders: reshape, gather, concat
17//! (plan #53). Other shape ops (narrow, transpose, expand) live
18//! on `GraphExt` in `infer.rs` since they need shape inference.
19
20use crate::{Graph, NodeId, Op, Shape};
21
22impl Graph {
23    /// Reshape.
24    pub fn reshape(&mut self, input: NodeId, new_shape: Vec<i64>, out_shape: Shape) -> NodeId {
25        self.push(Op::Reshape { new_shape }, vec![input], out_shape, None)
26    }
27
28    /// Gather (embedding lookup).
29    pub fn gather(&mut self, table: NodeId, indices: NodeId, axis: usize, shape: Shape) -> NodeId {
30        self.push(Op::Gather { axis }, vec![table, indices], shape, None)
31    }
32
33    /// Concatenate tensors along an axis.
34    pub fn concat(&mut self, inputs: Vec<NodeId>, axis: usize, shape: Shape) -> NodeId {
35        self.push(Op::Concat { axis }, inputs, shape, None)
36    }
37}
rlx_ir/ops/shape_ops.rs

rlx_ir/ops/
shape_ops.rs
