-
Notifications
You must be signed in to change notification settings - Fork 3
Description
Task description: Forward hooks and tape recorder implementation
-
Introduced a lightweight forward-hooks mechanism to record module forward passes without affecting normal execution when disabled.
-
Added
ForwardHooksinterface with two callbacks:onForwardBegin(module, input)andonForwardEnd(module, input, output). -
Extended the
ExecutionContextwith an optionalhooksproperty (nullmeans disabled), keeping zero cost for code paths when no hooks are provided. -
Plumbed hooks through concrete contexts:
DirectCpuExecutionContextandDefaultDataExecutionContextaccept and expose_hooksso callers can enable recording per run.
-
Implemented a simple
TapeRecorderthat collects entries (module id/name, input/output tensor specs: shape and dtype, and room for parameter refs) for each forward invocation. -
Added a small utility wrapper
withForwardHooks(ctx, module, input) { … }and used it to instrument forwards with minimal intrusion:- Integrated in
DualModulecomposition helpers so composed graphs are recorded as they execute. - Integrated in
Embedding.forward, ensuring index-based layers are recorded too.
- Integrated in
-
Kept traversal compatibility by relying on
ModuleNodefor id/name/children/params, so the tape can be correlated with the model structure. -
Wrote a test
ForwardHooksTapeTestthat:- Creates a
TapeRecorderand injects it intoDirectCpuExecutionContext. - Runs
Embedding.forwardon a small input and asserts exactly one module invocation was recorded with the expected module name and output shape in the tape.
- Creates a
-
Result: Developers can optionally attach a recorder for diagnostics, profiling, or future autodiff, with zero overhead when hooks are not supplied.