rlx_ir/ops/
shape_ops.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3// SPDX-License-Identifier: MIT OR Apache-2.0
4
5//! Shape-manipulation builders: reshape, gather, concat
6//! (plan #53). Other shape ops (narrow, transpose, expand) live
7//! on `GraphExt` in `infer.rs` since they need shape inference.
8
9use crate::op::ScatterNdReduction;
10use crate::{Graph, NodeId, Op, Shape};
11
12impl Graph {
13    /// Reshape.
14    pub fn reshape(&mut self, input: NodeId, new_shape: Vec<i64>, out_shape: Shape) -> NodeId {
15        self.push(Op::Reshape { new_shape }, vec![input], out_shape, None)
16    }
17
18    /// Gather (embedding lookup).
19    pub fn gather(&mut self, table: NodeId, indices: NodeId, axis: usize, shape: Shape) -> NodeId {
20        self.push(Op::Gather { axis }, vec![table, indices], shape, None)
21    }
22
23    /// ONNX ScatterND: copy `data`, write `updates` at multi-index locations.
24    pub fn scatter_nd(
25        &mut self,
26        data: NodeId,
27        indices: NodeId,
28        updates: NodeId,
29        reduction: ScatterNdReduction,
30    ) -> NodeId {
31        let shape = self.node(data).shape.clone();
32        self.push(
33            Op::ScatterNd { reduction },
34            vec![data, indices, updates],
35            shape,
36            None,
37        )
38    }
39
40    /// ONNX ScatterElements along `axis`.
41    pub fn scatter_elements(
42        &mut self,
43        data: NodeId,
44        indices: NodeId,
45        updates: NodeId,
46        axis: i32,
47        reduction: ScatterNdReduction,
48    ) -> NodeId {
49        let shape = self.node(data).shape.clone();
50        self.push(
51            Op::ScatterElements { axis, reduction },
52            vec![data, indices, updates],
53            shape,
54            None,
55        )
56    }
57
58    /// ONNX GatherND.
59    pub fn gather_nd(
60        &mut self,
61        data: NodeId,
62        indices: NodeId,
63        batch_dims: i32,
64        out_shape: Shape,
65    ) -> NodeId {
66        self.push(
67            Op::GatherNd { batch_dims },
68            vec![data, indices],
69            out_shape,
70            None,
71        )
72    }
73
74    /// ONNX GatherElements / take_along_axis — output shape = indices shape.
75    pub fn gather_elements(&mut self, data: NodeId, indices: NodeId, axis: i32) -> NodeId {
76        let shape = self.node(indices).shape.clone();
77        self.push(
78            Op::GatherElements { axis },
79            vec![data, indices],
80            shape,
81            None,
82        )
83    }
84
85    /// Concatenate tensors along an axis.
86    pub fn concat(&mut self, inputs: Vec<NodeId>, axis: usize, shape: Shape) -> NodeId {
87        self.push(Op::Concat { axis }, inputs, shape, None)
88    }
89
90    /// Reverse (flip) element order along each axis in `axes`. Output shape is
91    /// the same as the input; only the listed axes flip (batch-general).
92    pub fn reverse(&mut self, input: NodeId, axes: Vec<usize>) -> NodeId {
93        let shape = self.node(input).shape.clone();
94        self.push(Op::Reverse { axes }, vec![input], shape, None)
95    }
96}
rlx_ir/ops/shape_ops.rs

rlx_ir/ops/
shape_ops.rs
