Tensor hooks for gradients from each path in the autograd graph

username: akrentz6
cooked: <p>Hi all,</p>
<p>I’m working on a project where I need to manipulate gradients as they flow backward through each computational path in a model’s computation graph. Currently, I know that PyTorch’s <code>.register_hook</code> method allows you to manipulate gradients for a tensor during backpropagation, but this hook is only called after gradients from all incoming paths have been summed together.</p>
<p>I’d like to “hook in” before this summation step, so that I can observe or modify gradient contributions coming from each computational path (branch) individually, before they are combined at a given node in the graph.</p>
<p>I don’t think it’s possible to do this with PyTorch’s API, if anyone has suggestions that don’t involve modifying the C++ internals please let me know. I’m also wondering if anyone has any thoughts on if there would be demand for this kind of feature as an addition to the hook system - would it make sense as a feature request?</p>

username: ptrblck
cooked: <p>Have you tried using <a href="https://docs.pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module.register_full_backward_hook"><code>register_full_backward_hook</code></a>?<br>
The hook will give you access to all <code>grad_input</code> and <code>grad_output</code> entries.</p>

username: valerian.rey
cooked: <p>I’m also interested in this question. I made a minimal example of what happens with <code>register_hook</code>:</p>
<pre data-code-wrap="python"><code class="lang-python">a = torch.tensor(5., requires_grad=True)
b1 = 2 * a
b2 = 3 * a

def hook(grad):
    print(f"hook triggered with grad={grad}")

a.register_hook(hook)

torch.autograd.grad(outputs=[b1, b2], inputs=[a])
</code></pre>
<p>output: <code>hook triggered with grad=5.0</code></p>
<p>If I understand correctly you’re looking for a way to have a hook that would trigger twice: once with <code>grad=2.0</code> (i.e. the gradient of <code>b1</code> w.r.t. <code>a</code>) and once with <code>grad=3.0</code> (i.e. the gradient of <code>b2</code> w.r.t. <code>a</code>).</p>
<p>You can do that with a <a href="https://docs.pytorch.org/docs/stable/generated/torch.autograd.graph.Node.register_hook.html#torch.autograd.graph.Node.register_hook" rel="noopener nofollow ugc">post-hook on the autograd node</a> of <code>b1</code> and <code>b2</code>, rather than a pre-hook on <code>a</code>. This would look something like this:</p>
<pre data-code-wrap="python"><code class="lang-python">a = torch.tensor(5., requires_grad=True)
b1 = 2 * a
b2 = 3 * a

edge1 = get_gradient_edge(b1)
edge2 = get_gradient_edge(b2)

def hook(grad_inputs: tuple[Tensor], grad_outputs: tuple[Tensor]):
    print(f"hook triggered with grad={grad_inputs}")

edge1.node.register_hook(hook)
edge2.node.register_hook(hook)

torch.autograd.grad(outputs=[b1, b2], inputs=[a])
</code></pre>
<p>output:</p>
<pre><code class="lang-auto">hook triggered with grad=(tensor(3.), None)
hook triggered with grad=(tensor(2.), None)
</code></pre>
<p>I’m not sure why this gives tuple with also a <code>None</code> element, but at least we’ve got the desired two calls with the right gradients.</p>
<p>I think the solution proposed by <a class="mention" href="/u/ptrblck">@ptrblck</a> is similar to that, but in the case where the computations made from <code>a</code> are each within a <code>torch.module</code> (in which case a module post-hook, registered via <code>module.register_full_backward_hook</code> to each of the modules, would do the job).</p>
<p>I hope this helps <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=14" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"><br>
And please correct me if I misunderstood something!</p>

username: Hamza_Javaid
cooked: <p>Can you also try gradient edges?<br>
Writing a minimal example</p>
<pre><code class="lang-auto">import torch
from torch.autograd.graph import get_gradient_edge

def capture_individual_gradients(tensor, names=None):
    """
    Capture gradient contributions from each computational path
    before they're summed.
    
    Args:
        tensor: The tensor to monitor
        names: Optional list of names for each dependent operation
    
    Returns:
        dict: Maps operation names/indices to their gradient contributions
    """
    gradients = {}
    
    # Find all operations that depend on this tensor
    dependent_ops = []
    for op in tensor.grad_fn.next_functions if tensor.grad_fn else []:
        if op[0] is not None:
            dependent_ops.append(op)
    
    # Alternative: manually track operations
    def track_operation(output, name=None):
        edge = get_gradient_edge(output)
        
        def hook(grad_inputs, grad_outputs):
            # grad_inputs[0] contains the gradient flowing through this path
            grad_key = name if name else f"op_{len(gradients)}"
            gradients[grad_key] = grad_inputs[0].clone() if grad_inputs[0] is not None else None
            
        edge.node.register_hook(hook)
        return output
    
    return gradients, track_operation

# Example usage
a = torch.tensor(5., requires_grad=True)

# Set up gradient capture
gradients, track = capture_individual_gradients(a)

# Track each operation
b1 = track(2 * a, "multiply_by_2")
b2 = track(3 * a, "multiply_by_3")
b3 = track(a ** 2, "square")

# Compute gradients
loss = b1 + b2 + b3
loss.backward()

print("Individual gradients before summation:")
for name, grad in gradients.items():
    print(f"{name}: {grad}")
print(f"\nSummed gradient at 'a': {a.grad}")
</code></pre>

username: akrentz6
cooked: <p>Thanks all for the useful answers. I ended up using a modified version of <a class="mention" href="/u/hamza_javaid">@Hamza_Javaid</a>’s approach.</p>