Skip to content

[C] Add support for forward hooks #190

@michalharakal

Description

@michalharakal

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 ForwardHooks interface with two callbacks: onForwardBegin(module, input) and onForwardEnd(module, input, output).

  • Extended the ExecutionContext with an optional hooks property (null means disabled), keeping zero cost for code paths when no hooks are provided.

  • Plumbed hooks through concrete contexts:

    • DirectCpuExecutionContext and DefaultDataExecutionContext accept and expose _hooks so callers can enable recording per run.
  • Implemented a simple TapeRecorder that 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 DualModule composition helpers so composed graphs are recorded as they execute.
    • Integrated in Embedding.forward, ensuring index-based layers are recorded too.
  • Kept traversal compatibility by relying on ModuleNode for id/name/children/params, so the tape can be correlated with the model structure.

  • Wrote a test ForwardHooksTapeTest that:

    • Creates a TapeRecorder and injects it into DirectCpuExecutionContext.
    • Runs Embedding.forward on a small input and asserts exactly one module invocation was recorded with the expected module name and output shape in the tape.
  • Result: Developers can optionally attach a recorder for diagnostics, profiling, or future autodiff, with zero overhead when hooks are not supplied.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions