fandhe_ai/nn/
forward_hook.rs

1//! facade 版 forward hook ラッパー `ForwardHooked`（イシュー #2587・親 #2584）。
2//!
3//! 役割: facade 独自の [`crate::nn::Module`] 実装を包み、`forward` 完了直後に
4//! 観察専用 hook（PyTorch `register_forward_hook` 相当）を 1 回呼ぶ。autodiff 側
5//! `nn::ForwardHooked`（#2586）は autodiff `Module` 専用で、facade 利用者の層は
6//! 生の `fandhe_ai_autodiff::Tape` を名指せず直接は包めないため、本型が橋渡しをする
7//! （`docs/autodiff-forward-backward-hooks-design.md` §14.4 P5／P5′・§18）。
8//!
9//! 単一実装の再利用（P5′）: 内部では crate 内専用アダプタ `FacadeModuleAdapter` で
10//! inner を autodiff の `Module` へ包み、autodiff 側 `ForwardHooked` に hook の呼び出し・
11//! `ForwardHookCtx` の構築を任せる。facade 側は hook を直接呼ばず ctx も構築しない
12//! （ドリフトを避ける。`tests/api_surface.rs` が固定する）。`forward` 以外の `Module`
13//! メソッド（`named_parameters`・`state_dict` 等）は inner の facade `Module` メソッドへ
14//! 名前の接頭辞なしで透過委譲する。
15//!
16//! 契約:
17//! - hook は観察専用（戻り値は `Result<(), AutodiffError>`）。出力を書き換えられず、
18//!   `Err` は `forward` の `Err` として伝播する。inner の `forward` が `Err` なら hook は不発。
19//! - hook は構築時に 1 つだけ指定する（登録・解除 API は持たない）。
20//! - facade `Module` はホスト推論経路を持たない（`FacadeModuleAdapter` は
21//!   `supports_forward_host() == false`）ため、`compat::Sequential::predict` 等でも tape 経路で
22//!   1 forward につき 1 回発火する。
23//! - [`ForwardHookCtx::output_value`] は出力ノードを実体化するため、遅延融合の分割が変わりうる
24//!   （observer effect。値を読まない hook では出力は bit 一致。§5.3・§17.3）。
25//! - 包めるのは facade `nn::Module` 実装（利用者定義層・`nn::{ModuleList, Sequential,
26//!   ModuleDict}`）。autodiff 組み込み層の直接ラップは承認形の範囲外。
27//!
28//! 公開は `crate::nn::ForwardHooked` と `crate::nn::ForwardHookCtx`（autodiff からの純
29//! 再エクスポート。`nn/mod.rs`）のみで、`pub mod hooks` は設けない（P4）。
30
31use std::collections::HashMap;
32
33use fandhe_ai_autodiff::nn::ForwardHooked as AutodiffForwardHooked;
34use fandhe_ai_autodiff::nn::Module as AutodiffModule;
35
36use super::ForwardHookCtx;
37use crate::nn::module::{FacadeModuleAdapter, Module};
38use crate::{AutodiffError, TapeRef, Tensor, Var};
39
40/// facade `Module` を包み、`forward` 後に観察専用 hook を呼ぶラッパー。
41///
42/// 構築は [`Self::new`]、取り出しは [`Self::inner`]・[`Self::inner_mut`]・
43/// [`Self::into_inner`]。`Module` としては inner と同じ振る舞いを透過し、`forward` にだけ
44/// hook が加わる。hook の ctx（[`ForwardHookCtx`]）は入力・出力 shape と、必要時のみ出力値
45/// （`output_value`）を提供する。
46///
47/// ```
48/// use std::sync::Arc;
49/// use std::sync::atomic::{AtomicUsize, Ordering};
50/// use fandhe_ai::nn::{ForwardHooked, Module};
51/// use fandhe_ai::{AutodiffError, Tensor, TapeRef, Var};
52///
53/// struct Double;
54/// impl Module for Double {
55///     fn forward<'t>(
56///         &self,
57///         _tape: TapeRef<'t>,
58///         input: &Var<'t>,
59///     ) -> Result<Var<'t>, AutodiffError> {
60///         input.add(input)
61///     }
62/// }
63///
64/// let calls = Arc::new(AtomicUsize::new(0));
65/// let seen = Arc::clone(&calls);
66/// let hooked = ForwardHooked::new(Double, move |ctx| {
67///     assert_eq!(ctx.input_shape(), ctx.output_shape());
68///     seen.fetch_add(1, Ordering::SeqCst);
69///     Ok(())
70/// });
71/// let tape = fandhe_ai::tape();
72/// let x = tape.var(&Tensor::new(vec![1.0, 2.0], &[2]).unwrap());
73/// let y = hooked.forward(TapeRef::from(&tape), &x).unwrap();
74/// assert_eq!(y.to_tensor().host_slice().as_ref(), &[2.0, 4.0]);
75/// assert_eq!(calls.load(Ordering::SeqCst), 1);
76/// ```
77pub struct ForwardHooked<M: Module> {
78    inner: AutodiffForwardHooked<FacadeModuleAdapter<Box<M>>>,
79}
80
81impl<M: Module> ForwardHooked<M> {
82    /// `inner` を包み、`forward` 完了直後に `hook` を 1 回呼ぶラッパーを作る。
83    pub fn new<F>(inner: M, hook: F) -> Self
84    where
85        F: Fn(&ForwardHookCtx<'_>) -> Result<(), AutodiffError> + Send + Sync + 'static,
86    {
87        Self {
88            inner: AutodiffForwardHooked::new(FacadeModuleAdapter(Box::new(inner)), hook),
89        }
90    }
91
92    /// 内側の層への参照。
93    pub fn inner(&self) -> &M {
94        &self.inner.inner().0
95    }
96
97    /// 内側の層への可変参照。
98    pub fn inner_mut(&mut self) -> &mut M {
99        &mut self.inner.inner_mut().0
100    }
101
102    /// 包みを外して内側の層を取り出す（hook は破棄される）。
103    pub fn into_inner(self) -> M {
104        *self.inner.into_inner().0
105    }
106}
107
108impl<M: Module> Module for ForwardHooked<M> {
109    fn forward<'t>(&self, tape: TapeRef<'t>, input: &Var<'t>) -> Result<Var<'t>, AutodiffError> {
110        AutodiffModule::forward(&self.inner, tape.0, input)
111    }
112
113    fn named_parameters(&self) -> Vec<(String, &Tensor<f32>)> {
114        self.inner().named_parameters()
115    }
116
117    fn set_parameter(&mut self, name: &str, value: Tensor<f32>) -> Result<(), AutodiffError> {
118        self.inner_mut().set_parameter(name, value)
119    }
120
121    fn state_dict(&self) -> HashMap<String, Tensor<f32>> {
122        self.inner().state_dict()
123    }
124
125    fn load_state_dict(
126        &mut self,
127        state: HashMap<String, Tensor<f32>>,
128    ) -> Result<(), AutodiffError> {
129        self.inner_mut().load_state_dict(state)
130    }
131
132    fn set_training(&mut self, training: bool) {
133        self.inner_mut().set_training(training);
134    }
135
136    fn training(&self) -> bool {
137        self.inner().training()
138    }
139
140    fn set_requires_grad(&mut self, requires_grad: bool) -> Result<(), AutodiffError> {
141        self.inner_mut().set_requires_grad(requires_grad)
142    }
143
144    fn freeze(&mut self) -> Result<(), AutodiffError> {
145        self.inner_mut().freeze()
146    }
147
148    fn requires_grad(&self) -> bool {
149        self.inner().requires_grad()
150    }
151
152    fn children(&self) -> Vec<(String, &dyn Module)> {
153        self.inner().children()
154    }
155
156    fn children_mut(&mut self) -> Vec<(String, &mut dyn Module)> {
157        self.inner_mut().children_mut()
158    }
159
160    fn named_modules(&self) -> Vec<(String, &dyn Module)> {
161        self.inner().named_modules()
162    }
163
164    fn parameter_count(&self) -> usize {
165        self.inner().parameter_count()
166    }
167
168    fn type_name(&self) -> &'static str {
169        self.inner().type_name()
170    }
171}
fandhe_ai/nn/forward_hook.rs

fandhe_ai/nn/
forward_hook.rs
