rlx_tensor/
model.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//! A compiled, runnable [`Model`] built from an `rlx!` / graph-builder [`Graph`]:
6//! bind weights by name and run, without hand-writing the
7//! `Session` → `set_param` loop → `run` dance at every call site.
8
9use std::collections::HashSet;
10
11use rlx_ir::{Graph, Op};
12use rlx_runtime::{Device, Session};
13
14/// A compiled graph plus its parameter roster — a runnable model.
15///
16/// Bind weights with [`set`](Self::set) / [`load`](Self::load) *by name* (for an
17/// `rlx!` graph that name is the param's `@ "…"` key or, absent that, its
18/// binding ident), then [`run`](Self::run). Weight sources that carry extra
19/// tensors (a whole checkpoint) are fine — only the params this model declares
20/// are bound, and [`unbound`](Self::unbound) reports any it still needs.
21///
22/// ```ignore
23/// let mut m = Model::on(graph, Device::Cpu);
24/// m.load(weights.iter().map(|(k, v)| (k.as_str(), v.as_slice())));
25/// assert!(m.unbound().is_empty(), "missing weights: {:?}", m.unbound());
26/// let out = m.run(&[("input_ids", &ids)]);
27/// ```
28pub struct Model {
29    compiled: rlx_runtime::CompiledGraph,
30    params: HashSet<String>,
31    bound: HashSet<String>,
32}
33
34impl Model {
35    /// Compile on the fastest backend this build can run the graph on.
36    pub fn new(graph: Graph) -> Self {
37        let device = rlx_runtime::fastest_device_for(&graph);
38        Self::on(graph, device)
39    }
40
41    /// Compile on an explicit device.
42    pub fn on(graph: Graph, device: Device) -> Self {
43        let params = graph
44            .nodes()
45            .iter()
46            .filter_map(|n| match &n.op {
47                Op::Param { name } => Some(name.clone()),
48                _ => None,
49            })
50            .collect();
51        Self {
52            compiled: Session::new(device).compile(graph),
53            params,
54            bound: HashSet::new(),
55        }
56    }
57
58    /// Every parameter name the model expects.
59    pub fn param_names(&self) -> impl Iterator<Item = &str> {
60        self.params.iter().map(String::as_str)
61    }
62
63    /// Parameters not yet bound (in arbitrary order) — check this is empty after
64    /// [`load`](Self::load) to catch a missing / misnamed weight loudly.
65    pub fn unbound(&self) -> Vec<&str> {
66        self.params
67            .difference(&self.bound)
68            .map(String::as_str)
69            .collect()
70    }
71
72    /// Bind one parameter by name. A name the model doesn't declare is ignored
73    /// (so binding straight from a superset checkpoint is safe).
74    pub fn set(&mut self, name: &str, data: &[f32]) -> &mut Self {
75        if self.params.contains(name) {
76            self.compiled.set_param(name, data);
77            self.bound.insert(name.to_string());
78        }
79        self
80    }
81
82    /// Bind every parameter the model declares from a `name → data` source
83    /// (a `HashMap`, a slice of pairs, …); extra entries are ignored.
84    pub fn load<'a, I>(&mut self, weights: I) -> &mut Self
85    where
86        I: IntoIterator<Item = (&'a str, &'a [f32])>,
87    {
88        for (name, data) in weights {
89            self.set(name, data);
90        }
91        self
92    }
93
94    /// Run with named inputs, returning one `Vec<f32>` per graph output.
95    pub fn run(&mut self, inputs: &[(&str, &[f32])]) -> Vec<Vec<f32>> {
96        self.compiled.run(inputs)
97    }
98}
rlx_tensor/model.rs

rlx_tensor/
model.rs
