Skip to main content

optirs_learned/lstm/
introspection.rs

1//! Read-only views of the LSTM controller's learned weights.
2//!
3//! Initialization statistics (finding F64) and gradient-flow diagnostics can only
4//! be asserted from outside the crate if the weight matrices are reachable, and
5//! `LSTMNetwork`/`LSTMLayer`/`OutputProjection` keep all of theirs private.
6//! These accessors hand out immutable borrows only — nothing here lets a caller
7//! mutate a weight, so the invariants the forward pass relies on (matching
8//! `weights.ncols()` and input width, four gate blocks stacked along the rows)
9//! stay under the owning type's control.
10//!
11//! Kept in its own file so `lstm.rs` stays under the 2000-line cap.
12
13use scirs2_core::ndarray::{Array1, Array2};
14use scirs2_core::numeric::Float;
15use std::fmt::Debug;
16
17use super::{AttentionMechanism, LSTMLayer, LSTMNetwork, LSTMOptimizer, OutputProjection};
18
19impl<T: Float + Debug + Send + Sync + 'static> OutputProjection<T> {
20    /// The projection matrix, shaped `(output_size, input_size)`.
21    pub fn weight_snapshot(&self) -> &Array2<T> {
22        &self.weights
23    }
24
25    /// The projection bias, length `output_size`.
26    pub fn bias_snapshot(&self) -> &Array1<T> {
27        &self.bias
28    }
29}
30
31impl<T: Float + Debug + Send + Sync + 'static> LSTMLayer<T> {
32    /// Input-to-hidden weights, shaped `(4 · hidden_size, input_size)`.
33    ///
34    /// The four gate blocks (input, forget, cell, output) are stacked along the
35    /// rows, so the *fan-out* of each gate is `hidden_size`, not `4 ·
36    /// hidden_size` — which is what the Glorot limit has to be computed from.
37    pub fn weight_ih_snapshot(&self) -> &Array2<T> {
38        &self.weight_ih
39    }
40
41    /// Hidden-to-hidden weights, shaped `(4 · hidden_size, hidden_size)`.
42    pub fn weight_hh_snapshot(&self) -> &Array2<T> {
43        &self.weight_hh
44    }
45}
46
47impl<T: Float + Debug + Send + Sync + 'static> AttentionMechanism<T> {
48    /// The four attention projections in `(query, key, value, output)` order,
49    /// each shaped `(hidden_size, hidden_size)`.
50    pub fn projection_snapshots(&self) -> (&Array2<T>, &Array2<T>, &Array2<T>, &Array2<T>) {
51        (
52            &self.query_proj,
53            &self.key_proj,
54            &self.value_proj,
55            &self.output_proj,
56        )
57    }
58}
59
60impl<T: Float + Debug + Send + Sync + 'static> LSTMNetwork<T> {
61    /// Number of stacked LSTM layers.
62    pub fn layer_count(&self) -> usize {
63        self.layers.len()
64    }
65
66    /// Borrow the `index`-th LSTM layer, or `None` if it does not exist.
67    pub fn layer(&self, index: usize) -> Option<&LSTMLayer<T>> {
68        self.layers.get(index)
69    }
70
71    /// Borrow the output projection.
72    pub fn output_projection_snapshot(&self) -> &OutputProjection<T> {
73        &self.output_projection
74    }
75
76    /// Borrow the attention mechanism, if the configuration enabled one.
77    pub fn attention_snapshot(&self) -> Option<&AttentionMechanism<T>> {
78        self.attention.as_ref()
79    }
80}
81
82impl<T: Float + Debug + Send + Sync + 'static> LSTMOptimizer<T> {
83    /// `(weight_ih, weight_hh)` of the first LSTM layer, or `None` when the
84    /// controller has no layers.
85    pub fn first_layer_weight_snapshot(&self) -> Option<(&Array2<T>, &Array2<T>)> {
86        self.lstm_network
87            .layer(0)
88            .map(|layer| (layer.weight_ih_snapshot(), layer.weight_hh_snapshot()))
89    }
90
91    /// Borrow the controller network for read-only inspection.
92    pub fn network_snapshot(&self) -> &LSTMNetwork<T> {
93        &self.lstm_network
94    }
95}