Autograd graph traversal

username: valerian.rey
cooked: <h2><a name="p-460509-context-1" class="anchor" href="#p-460509-context-1"></a>Context</h2>
<p>I’m developing a library on top of PyTorch for Jacobian descent, called <a href="https://github.com/TorchJD/torchjd" rel="noopener nofollow ugc">torchjd</a>. I have a function <code>torchjd.backward</code>, that I want to be quite similar to <code>torch.backward</code> in its interface. Basically, the difference is that <code>torch.backward</code> computes the sum of gradients of given tensors with respect to graph leaves, while <code>torchjd.backward</code> can do more complex aggregations than summing the gradients.</p>
<h2><a name="p-460509-problem-2" class="anchor" href="#p-460509-problem-2"></a>Problem</h2>
<p>I don’t know how to cleanly find the default <code>inputs</code> with respect to which <code>tensors</code> need to<br>
be differentiated. If those are specified manually, my function works, but similarly to <code>torch.backward</code>, I’d like the <code>inputs</code> parameter to default to the leaf tensors that were used to obtain the <code>tensors</code> parameter.<br>
In the case of <code>torch.backward</code>, these leaf tensors are found through the C++ code. In my case, I would like to keep the library purely python, so I think that I need to find those tensors by traversing the backward computation graph.</p>
<h2><a name="p-460509-question-3" class="anchor" href="#p-460509-question-3"></a>Question</h2>
<p>Given a list of tensors, how can I find all leaf tensors requiring grad that were used to compute these tensors?</p>
<p>Currently, our proposed implementation looks like this</p>
<pre data-code-wrap="python"><code class="lang-python">def _get_leaves_of_autograd_graph(roots: list[Tensor]) -&gt; set[Tensor]:
    """
    Gets the leafs of the autograd graph of all tensors in `roots`.

    :param roots: Roots of the autograd graphs.
    """
    nodes_to_traverse = [tensor.grad_fn for tensor in roots]
    leaves = set()

    while nodes_to_traverse:
        current_node = nodes_to_traverse.pop()

        if hasattr(current_node, "variable"):
            leaves.add(current_node.variable)
        else:
            nodes_to_traverse += [child[0] for child in current_node.next_functions if child[0] is not None]

    return leaves
</code></pre>
<p>I’m really worried about the <code>hasattr(current_node, "variable")</code> part, since I don’t think there is any guarantee that the variable of a node will be stored in the <code>variable</code> field.</p>
<p>Does anyone know how this implementation could be improved? At the moment, I don’t think it is solid enough to be integrated into our library.</p>
<p>Also, shouldn’t PyTorch itself provide an easier way to find such values, or a more standard way to traverse the AG graph?</p>
<p>Thanks for reading this long question. Any help would be really appreciated!</p>

username: soulitzer
cooked: <blockquote>
<p>nodes_to_traverse = [tensor.grad_fn for tensor in roots]</p>
</blockquote>
<p>one edge case to note is that if tensor itself were a leaf, grad_fn would be None</p>
<blockquote>
<p>hasattr(current_node, “variable”)</p>
</blockquote>
<p>I don’t expect this to change, should be fine to use it.</p>
<blockquote>
<p>Also, shouldn’t PyTorch itself provide an easier way to find such values, or a more standard way to traverse the AG graph?</p>
</blockquote>
<p>Sounds reasonable -</p>
<p>We already have this generator that is used a couple places in the code, and we can probably expose it.</p>
<pre><code class="lang-auto">def iter_graph(roots: List[Node]) -&gt; Iterator[Node]:
        if not roots:
            return
        seen: Set[Node] = set()
        q: Deque[Node] = deque()
        for node in roots:
            if node is not None and node not in seen::
                seen.add(node)
                q.append(node)

        while q:
            node = q.popleft()
            for fn, _ in node.next_functions:
                if fn in seen or fn is None:
                    continue
                seen.add(fn)
                q.append(fn)

            yield node
</code></pre>
<p>This seems reasonable to have, we’d accept a PR implementing this</p>

username: valerian.rey
cooked: <p>Thanks a lot for the answer! I will try to make a PR in the coming days. I will tag you on it when it’s ready, so that we can keep discussing there.</p>