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}