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::op::ScatterNdReduction;
21use crate::{Graph, NodeId, Op, Shape};
22
23impl Graph {
24    /// Reshape.
25    pub fn reshape(&mut self, input: NodeId, new_shape: Vec<i64>, out_shape: Shape) -> NodeId {
26        self.push(Op::Reshape { new_shape }, vec![input], out_shape, None)
27    }
28
29    /// Gather (embedding lookup).
30    pub fn gather(&mut self, table: NodeId, indices: NodeId, axis: usize, shape: Shape) -> NodeId {
31        self.push(Op::Gather { axis }, vec![table, indices], shape, None)
32    }
33
34    /// ONNX ScatterND: copy `data`, write `updates` at multi-index locations.
35    pub fn scatter_nd(
36        &mut self,
37        data: NodeId,
38        indices: NodeId,
39        updates: NodeId,
40        reduction: ScatterNdReduction,
41    ) -> NodeId {
42        let shape = self.node(data).shape.clone();
43        self.push(
44            Op::ScatterNd { reduction },
45            vec![data, indices, updates],
46            shape,
47            None,
48        )
49    }
50
51    /// ONNX ScatterElements along `axis`.
52    pub fn scatter_elements(
53        &mut self,
54        data: NodeId,
55        indices: NodeId,
56        updates: NodeId,
57        axis: i32,
58        reduction: ScatterNdReduction,
59    ) -> NodeId {
60        let shape = self.node(data).shape.clone();
61        self.push(
62            Op::ScatterElements { axis, reduction },
63            vec![data, indices, updates],
64            shape,
65            None,
66        )
67    }
68
69    /// ONNX GatherND.
70    pub fn gather_nd(
71        &mut self,
72        data: NodeId,
73        indices: NodeId,
74        batch_dims: i32,
75        out_shape: Shape,
76    ) -> NodeId {
77        self.push(
78            Op::GatherNd { batch_dims },
79            vec![data, indices],
80            out_shape,
81            None,
82        )
83    }
84
85    /// ONNX GatherElements / take_along_axis — output shape = indices shape.
86    pub fn gather_elements(&mut self, data: NodeId, indices: NodeId, axis: i32) -> NodeId {
87        let shape = self.node(indices).shape.clone();
88        self.push(
89            Op::GatherElements { axis },
90            vec![data, indices],
91            shape,
92            None,
93        )
94    }
95
96    /// Concatenate tensors along an axis.
97    pub fn concat(&mut self, inputs: Vec<NodeId>, axis: usize, shape: Shape) -> NodeId {
98        self.push(Op::Concat { axis }, inputs, shape, None)
99    }
100
101    /// Reverse (flip) element order along each axis in `axes`. Output shape is
102    /// the same as the input; only the listed axes flip (batch-general).
103    pub fn reverse(&mut self, input: NodeId, axes: Vec<usize>) -> NodeId {
104        let shape = self.node(input).shape.clone();
105        self.push(Op::Reverse { axes }, vec![input], shape, None)
106    }
107}
rlx_ir/ops/shape_ops.rs

rlx_ir/ops/
shape_ops.rs
