xlog_neural/
handle.rs

1//! Network handle for managing PyTorch modules.
2//!
3//! Each registered neural network is represented by a `NetworkHandle` which holds:
4//! - The PyTorch module (nn.Module) - when `python` feature is enabled
5//! - Optional optimizer for training
6//! - Optional learning rate scheduler
7//! - Configuration flags for batching, caching, etc.
8
9#[cfg(feature = "python")]
10use pyo3::{Py, PyAny};
11
12/// Handle to a registered neural network.
13///
14/// This struct holds the PyTorch module and associated training state.
15/// When the `python` feature is enabled, it can hold PyO3 `Py<PyAny>` references.
16#[derive(Debug)]
17pub struct NetworkHandle {
18    /// Unique name identifying this network
19    pub name: String,
20
21    /// The PyTorch nn.Module (set via Python API)
22    /// Only available with the `python` feature
23    #[cfg(feature = "python")]
24    pub module: Option<Py<PyAny>>,
25
26    /// The optimizer for training (e.g., Adam, SGD)
27    /// Only available with the `python` feature
28    #[cfg(feature = "python")]
29    pub optimizer: Option<Py<PyAny>>,
30
31    /// Learning rate scheduler
32    /// Only available with the `python` feature
33    #[cfg(feature = "python")]
34    pub scheduler: Option<Py<PyAny>>,
35
36    /// Whether to batch inputs for efficient GPU processing
37    pub batching: bool,
38
39    /// Top-k sampling: if Some(k), only consider top k outputs
40    pub k: Option<usize>,
41
42    /// Deterministic mode: use argmax instead of sampling
43    pub det: bool,
44
45    /// Whether the network is in training mode
46    pub train_mode: bool,
47
48    /// Whether output caching is enabled
49    pub cache_enabled: bool,
50
51    /// Maximum number of cached outputs
52    pub cache_size: usize,
53
54    /// Registration-time arity, as passed to `register_network` (see
55    /// `NetworkConfig::arity`). Carried through unchanged so it can be read
56    /// back after registration; the handle does not interpret it.
57    pub arity: Option<usize>,
58
59    /// Registration-time per-argument catalog sort ids (see
60    /// `NetworkConfig::arg_sorts`). Carried through unchanged, like `arity`.
61    pub arg_sorts: Option<Vec<i64>>,
62
63    /// Registration-time artifact content hash (see `NetworkConfig::artifact_hash`).
64    /// Carried through unchanged, like `arity`.
65    pub artifact_hash: Option<String>,
66}
67
68impl NetworkHandle {
69    /// Create a new network handle with the given name and default settings.
70    pub fn new(name: String) -> Self {
71        Self {
72            name,
73            #[cfg(feature = "python")]
74            module: None,
75            #[cfg(feature = "python")]
76            optimizer: None,
77            #[cfg(feature = "python")]
78            scheduler: None,
79            batching: true,
80            k: None,
81            det: false,
82            train_mode: false,
83            cache_enabled: true,
84            cache_size: 10000,
85            arity: None,
86            arg_sorts: None,
87            artifact_hash: None,
88        }
89    }
90
91    /// Create a handle from a configuration.
92    pub fn from_config(config: &crate::NetworkConfig) -> Self {
93        Self {
94            name: config.name.clone(),
95            #[cfg(feature = "python")]
96            module: None,
97            #[cfg(feature = "python")]
98            optimizer: None,
99            #[cfg(feature = "python")]
100            scheduler: None,
101            batching: config.batching,
102            k: config.k,
103            det: config.det,
104            train_mode: false,
105            cache_enabled: config.cache_enabled,
106            cache_size: config.cache_size,
107            arity: config.arity,
108            arg_sorts: config.arg_sorts.clone(),
109            artifact_hash: config.artifact_hash.clone(),
110        }
111    }
112
113    /// Check if the PyTorch module has been set.
114    #[cfg(feature = "python")]
115    pub fn has_module(&self) -> bool {
116        self.module.is_some()
117    }
118
119    /// Check if the PyTorch module has been set.
120    /// Without Python feature, always returns false.
121    #[cfg(not(feature = "python"))]
122    /// Report whether a Python module/tensor handle is attached.
123    pub fn has_module(&self) -> bool {
124        false
125    }
126
127    /// Check if an optimizer has been configured.
128    #[cfg(feature = "python")]
129    pub fn has_optimizer(&self) -> bool {
130        self.optimizer.is_some()
131    }
132
133    /// Check if an optimizer has been configured.
134    /// Without Python feature, always returns false.
135    #[cfg(not(feature = "python"))]
136    pub fn has_optimizer(&self) -> bool {
137        false
138    }
139
140    /// Check if a scheduler has been configured.
141    #[cfg(feature = "python")]
142    pub fn has_scheduler(&self) -> bool {
143        self.scheduler.is_some()
144    }
145
146    /// Check if a scheduler has been configured.
147    /// Without Python feature, always returns false.
148    #[cfg(not(feature = "python"))]
149    pub fn has_scheduler(&self) -> bool {
150        false
151    }
152
153    /// Set the PyTorch module.
154    #[cfg(feature = "python")]
155    pub fn set_module(&mut self, module: Py<PyAny>) {
156        self.module = Some(module);
157    }
158
159    /// Set the optimizer.
160    #[cfg(feature = "python")]
161    pub fn set_optimizer(&mut self, optimizer: Py<PyAny>) {
162        self.optimizer = Some(optimizer);
163    }
164
165    /// Set the learning rate scheduler.
166    #[cfg(feature = "python")]
167    pub fn set_scheduler(&mut self, scheduler: Py<PyAny>) {
168        self.scheduler = Some(scheduler);
169    }
170
171    /// Get a reference to the PyTorch module.
172    #[cfg(feature = "python")]
173    pub fn module(&self) -> Option<&Py<PyAny>> {
174        self.module.as_ref()
175    }
176
177    /// Get a reference to the optimizer.
178    #[cfg(feature = "python")]
179    pub fn optimizer(&self) -> Option<&Py<PyAny>> {
180        self.optimizer.as_ref()
181    }
182
183    /// Get a reference to the scheduler.
184    #[cfg(feature = "python")]
185    pub fn scheduler(&self) -> Option<&Py<PyAny>> {
186        self.scheduler.as_ref()
187    }
188
189    /// Clear the module and training state.
190    #[cfg(feature = "python")]
191    pub fn clear(&mut self) {
192        self.module = None;
193        self.optimizer = None;
194        self.scheduler = None;
195    }
196
197    /// Clear the module and training state.
198    /// Without Python feature, this is a no-op.
199    #[cfg(not(feature = "python"))]
200    pub fn clear(&mut self) {
201        // No-op without Python feature
202    }
203}
204
205/// Handle to a registered embedding module.
206///
207/// Wraps either a trainable `nn.Embedding` or a frozen `torch.Tensor`.
208/// Created via `CompiledProgram.register_embedding()` in Python.
209#[derive(Debug)]
210pub struct EmbeddingHandle {
211    /// Unique name matching the nn() declaration
212    pub name: String,
213
214    /// The PyTorch nn.Embedding or tensor
215    #[cfg(feature = "python")]
216    pub module: Option<Py<PyAny>>,
217
218    /// Whether gradients flow through this embedding
219    pub trainable: bool,
220
221    /// Embedding vector dimension (second axis of weight matrix)
222    pub dim: usize,
223
224    /// Number of embedding entries (first axis of weight matrix)
225    pub vocab_size: usize,
226}
227
228impl EmbeddingHandle {
229    /// Create a new embedding handle.
230    pub fn new(name: String, trainable: bool, dim: usize, vocab_size: usize) -> Self {
231        Self {
232            name,
233            #[cfg(feature = "python")]
234            module: None,
235            trainable,
236            dim,
237            vocab_size,
238        }
239    }
240
241    /// Check if the PyTorch module/tensor has been set.
242    #[cfg(feature = "python")]
243    pub fn has_module(&self) -> bool {
244        self.module.is_some()
245    }
246
247    /// Check if the PyTorch module/tensor has been set.
248    /// Without Python feature, always returns false.
249    #[cfg(not(feature = "python"))]
250    pub fn has_module(&self) -> bool {
251        false
252    }
253
254    /// Set the PyTorch module/tensor.
255    #[cfg(feature = "python")]
256    pub fn set_module(&mut self, module: Py<PyAny>) {
257        self.module = Some(module);
258    }
259
260    /// Get a reference to the PyTorch module/tensor.
261    #[cfg(feature = "python")]
262    pub fn module(&self) -> Option<&Py<PyAny>> {
263        self.module.as_ref()
264    }
265}
266
267impl Default for NetworkHandle {
268    fn default() -> Self {
269        Self::new(String::new())
270    }
271}
272
273#[cfg(test)]
274mod tests {
275    use super::*;
276
277    #[test]
278    fn test_handle_new() {
279        let handle = NetworkHandle::new("test".to_string());
280        assert_eq!(handle.name, "test");
281        assert!(!handle.has_module());
282        assert!(!handle.has_optimizer());
283        assert!(handle.batching);
284        assert!(!handle.train_mode);
285    }
286
287    #[test]
288    fn test_handle_from_config() {
289        let config = crate::NetworkConfig {
290            name: "configured".to_string(),
291            batching: false,
292            k: Some(5),
293            det: true,
294            cache_enabled: false,
295            cache_size: 500,
296            arity: None,
297            arg_sorts: None,
298            artifact_hash: None,
299        };
300
301        let handle = NetworkHandle::from_config(&config);
302        assert_eq!(handle.name, "configured");
303        assert!(!handle.batching);
304        assert_eq!(handle.k, Some(5));
305        assert!(handle.det);
306        assert!(!handle.cache_enabled);
307        assert_eq!(handle.cache_size, 500);
308        assert!(handle.arity.is_none());
309        assert!(handle.arg_sorts.is_none());
310        assert!(handle.artifact_hash.is_none());
311    }
312
313    #[test]
314    fn test_handle_from_config_carries_registration_metadata() {
315        // The handle is what survives after `register()` consumes the config
316        // (see NetworkRegistry::register), so registration metadata that
317        // isn't copied here is lost the moment registration completes.
318        let mut config = crate::NetworkConfig::default("configured");
319        config.arity = Some(2);
320        config.arg_sorts = Some(vec![0, 1]);
321        config.artifact_hash = Some("deadbeef".to_string());
322
323        let handle = NetworkHandle::from_config(&config);
324        assert_eq!(handle.arity, Some(2));
325        assert_eq!(handle.arg_sorts, Some(vec![0, 1]));
326        assert_eq!(handle.artifact_hash, Some("deadbeef".to_string()));
327    }
328
329    #[test]
330    fn test_embedding_handle_new() {
331        let handle = EmbeddingHandle::new("test_embed".to_string(), true, 64, 1000);
332        assert_eq!(handle.name, "test_embed");
333        assert!(handle.trainable);
334        assert_eq!(handle.dim, 64);
335        assert_eq!(handle.vocab_size, 1000);
336        assert!(!handle.has_module());
337    }
338
339    #[test]
340    fn test_embedding_handle_frozen() {
341        let handle = EmbeddingHandle::new("frozen".to_string(), false, 128, 500);
342        assert!(!handle.trainable);
343        assert_eq!(handle.dim, 128);
344        assert_eq!(handle.vocab_size, 500);
345    }
346}
xlog_neural/handle.rs

xlog_neural/
handle.rs
