rlx_tensor/
scope.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3
4//! Graph builder scope.
5
6use rlx_ir::{DType, Dim, Graph, GraphExt, NodeId, Shape};
7
8use crate::Tensor;
9use crate::handle::GraphHandle;
10use crate::scalar::promote_scalar;
11
12/// Owns a traced graph while you build symbolic tensors.
13#[derive(Debug)]
14pub struct GraphScope {
15    handle: GraphHandle,
16    /// Monotonic counter for fresh dynamic-dim symbols.
17    next_dyn: u32,
18}
19
20impl GraphScope {
21    pub fn new(name: impl Into<String>) -> Self {
22        Self {
23            handle: GraphHandle::new(Graph::new(name)),
24            next_dyn: 0,
25        }
26    }
27
28    pub fn graph(&self) -> Graph {
29        self.handle.with_graph(|g| g.clone())
30    }
31
32    /// Fresh dynamic-dimension symbol (for variable batch / sequence length).
33    pub fn fresh_symbol(&mut self) -> u32 {
34        let s = self.next_dyn;
35        self.next_dyn += 1;
36        s
37    }
38
39    pub fn input(&mut self, name: impl Into<String>, shape: impl Into<Shape>) -> Tensor {
40        let id = self.handle.with_graph(|g| g.input(name, shape.into()));
41        Tensor::new(self.handle.clone(), id)
42    }
43
44    /// Input with a caller-chosen dynamic symbol on one axis.
45    pub fn input_dyn(&mut self, name: impl Into<String>, dims: &[Dim], dtype: DType) -> Tensor {
46        let shape = Shape::from_dims(dims, dtype);
47        self.input(name, shape)
48    }
49
50    pub fn param(&mut self, name: impl Into<String>, shape: impl Into<Shape>) -> Tensor {
51        let id = self.handle.with_graph(|g| g.param(name, shape.into()));
52        Tensor::new(self.handle.clone(), id)
53    }
54
55    /// Rank-0 literal broadcastable in binary ops.
56    pub fn scalar(&mut self, value: f64) -> Tensor {
57        let id = self
58            .handle
59            .with_graph(|g| promote_scalar(g, value.into(), DType::F32));
60        Tensor::new(self.handle.clone(), id)
61    }
62
63    pub fn constant(&mut self, value: f64, dtype: DType) -> Tensor {
64        let id = self.handle.with_graph(|g| g.constant(value, dtype));
65        Tensor::new(self.handle.clone(), id)
66    }
67
68    /// Stack existing tensors along `axis` (zero-copy at IR level).
69    pub fn cat(&mut self, tensors: &[&Tensor], axis: usize) -> Tensor {
70        let ids: Vec<NodeId> = tensors.iter().map(|t| t.id).collect();
71        let id = self.handle.with_graph(|g| g.concat_(ids, axis));
72        Tensor::new(self.handle.clone(), id)
73    }
74
75    pub fn set_outputs(&mut self, outputs: impl IntoIterator<Item = impl Into<NodeId>>) {
76        let ids: Vec<NodeId> = outputs.into_iter().map(Into::into).collect();
77        self.handle.with_graph(|g| g.set_outputs(ids));
78    }
79
80    pub fn finish(self) -> Graph {
81        self.handle.with_graph(|g| g.clone())
82    }
83}
84
85/// Build a single-output graph from a closure.
86///
87/// ```rust
88/// use rlx_tensor::{graph, shape};
89///
90/// let g = graph("mlp", |g| {
91///     let x = g.input("x", shape![2, 4]);
92///     let w = g.param("w", shape![4, 3]);
93///     let b = g.param("b", shape![3]);
94///     (&x.matmul(&w) + &b).gelu() * 2.0
95/// });
96/// assert_eq!(g.outputs.len(), 1);
97/// ```
98pub fn graph<F>(name: impl Into<String>, build: F) -> Graph
99where
100    F: FnOnce(&mut GraphScope) -> Tensor,
101{
102    let mut scope = GraphScope::new(name);
103    let out = build(&mut scope);
104    scope.set_outputs([out]);
105    scope.finish()
106}
107
108/// Build a graph and return extra metadata from the closure.
109pub fn graph_with<F, R>(name: impl Into<String>, build: F) -> (Graph, R)
110where
111    F: FnOnce(&mut GraphScope) -> R,
112{
113    let mut scope = GraphScope::new(name);
114    let result = build(&mut scope);
115    (scope.finish(), result)
116}
rlx_tensor/scope.rs

rlx_tensor/
scope.rs
