class

ModuleBackwardHookState

ModuleBackwardHookState

source edit

Shared state shuttled between the input-side and output-side hook barriers of a single nn.Module invocation.

A Module with registered backward hooks is wrapped at forward time with two barriers — ModuleOutputHookNode on each output edge and ModuleInputHookNode on each gradient-receiving input edge. Both barriers hold a shared_ptr to the same ModuleBackwardHookState so that the gradients arriving at the outputs (from downstream nodes) are available to the input-side hook callbacks, and any replacement gradients returned by the Python full-backward hook flow back through the inputs.

Attributes

pre_runnerpy::object

Python callable that runs the pre-backward hooks. Receives the output-gradient tuple, may return a tuple replacing it.

full_runnerpy::object

Python callable that runs the full-backward hooks. Receives both the input-gradient tuple and the output-gradient tuple, may return a tuple replacing the input gradients.

n_inputsstd::size_t

Number of original Python-side input positions on the module call (some may have requires_grad=False and therefore no edge).

n_outputsstd::size_t

Number of outputs the module produced.

pre_hooks_ranbool

One-shot guard so the pre-backward hooks fire exactly once per backward pass even when the barrier is visited from multiple edges.

full_hooks_ranbool

Same guard for the full-backward hook batch.

input_arg_indicesstd::vector<std::uint32_t>

Maps each edge slot in grad_inputs back to its original positional argument index in the Python signature.

input_metasstd::vector<ModuleHookTensorMeta>

Metadata for each gradient-receiving input edge, parallel to grad_inputs.

grad_inputsstd::vector<std::optional<Storage>>

Per-edge incoming gradient buffers populated by the input-side barrier. std::nullopt means no gradient arrived for that slot.

output_metasstd::vector<ModuleHookTensorMeta>

Metadata for every output position (size() == n_outputs).

grad_outputsstd::vector<std::optional<Storage>>

Gradient buffers accumulated at the output barrier, indexed by output position.

output_edge_indicesstd::vector<std::uint32_t>

For each edge of the output barrier, the output-position index it corresponds to in grad_outputs.

Notes

The two py::object runners are the Python-side dispatchers that fan each barrier event out to every registered hook on the module. They are held as opaque py::object (the GIL is acquired only when they are actually invoked inside apply_barrier).

Constructors

ctor