Sparse Jacobian computation

username: yaoyuh
cooked: <p>Hi PyTorch people,</p>
<p>I need to compute the Jacobian of a multi-variant vector-valued function w.r.t. its inputs. The problem I have is the dimension of the Jacobian is too large and it consumes a lot of GPU memory. What I’d like to do is to augment the current PyTorch implementation of <a href="https://github.com/pytorch/pytorch/blob/5ffda02388f1a1a3c83d8e6676ec1c7019c5ecd1/torch/autograd/functional.py#L487" rel="noopener nofollow ugc">torch.autograd.functional.jacobian</a> and somehow return the Jacobian represented as a block sparse matrix. However, I run into an immediate problem that for every individual output element, I need to know whether or not an input element has participated in the computation. Normally, PyTorch gives us zero gradients if an input has not been used for computing a certain output element.</p>
<p>For example, if I run the following sample code,</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
from torch.autograd.functional import jacobian

def func(a0, a1, a2):
    a0 = a0 * 0
    a1 = a1 * 2
    a2 = a2 * 4
    e = torch.cat( (a0, a1, a2) )
    return e

def func_single(g):
    return func( g[:2], g[2:4], g[4:6] )

def test_jacobian_participation():
    print() # Create an empty line.

    a0 = torch.rand((2,), requires_grad=True)
    a1 = torch.rand((2,), requires_grad=True)
    a2 = torch.rand((2,), requires_grad=True)

    g = torch.zeros((6,))
    g[:2]  = a0
    g[2:4] = a1
    g[4:6] = a2

    j = jacobian(func_single, g)
    print(f'j = \n{j}')

def test_jacobian_participation_tuple():
    print() # Create an empty line.

    a0 = torch.rand((2,), requires_grad=True)
    a1 = torch.rand((2,), requires_grad=True)
    a2 = torch.rand((2,), requires_grad=True)

    j_tuple = jacobian( func, (a0, a1, a2) )

    for i, j in enumerate(j_tuple):
        print(f'j_{i} = \n{j}')

if __name__ == '__main__':
    test_jacobian_participation()
    test_jacobian_participation_tuple()
</code></pre>
<p>I got the following in the terminal</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/6/a/6a87462c776b4baeaa62d67624ae2f9db166ccd8.png" data-download-href="https://discuss.pytorch.org/uploads/default/6a87462c776b4baeaa62d67624ae2f9db166ccd8" title="Screenshot from 2022-09-02 17-07-28"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/6/a/6a87462c776b4baeaa62d67624ae2f9db166ccd8_2_316x500.png" alt="Screenshot from 2022-09-02 17-07-28" data-base62-sha1="fcoqZo0U9v3YlR0NyKs3Pn8psCs" width="316" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/6/a/6a87462c776b4baeaa62d67624ae2f9db166ccd8_2_316x500.png, https://discuss.pytorch.org/uploads/default/original/3X/6/a/6a87462c776b4baeaa62d67624ae2f9db166ccd8.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/6/a/6a87462c776b4baeaa62d67624ae2f9db166ccd8.png 2x" data-dominant-color="36182C"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screenshot from 2022-09-02 17-07-28</span><span class="informations">350×553 48.5 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>During the call to the jacobian function, PyTorch does the computation in a row-based manner. That is the values for the Jacobian are computed <a href="https://github.com/pytorch/pytorch/blob/5ffda02388f1a1a3c83d8e6676ec1c7019c5ecd1/torch/autograd/functional.py#L672" rel="noopener nofollow ugc">row-by-row</a>, as shown in the red dashed frames in the above figure. I am planning to modify this such that I can accumulate row results until I have enough rows for a single output object (such as 2 rows for a 2-vector representing the XY location of a 2D point, green and blue boxes in the above figure). One problem I have is that when I see one row or several rows of Jacobian values, how do I tell if a subset of zeros are there because of</p>
<ul>
<li>The associated input element does not participate in the computation of the output element (green boxes, should <strong>not</strong> go into the sparse block matrix).</li>
<li>The associated input element participate in the computation but got multiplied by zero (blue boxes, <strong>should go</strong> into the sparse block matrix).</li>
</ul>
<p>Any suggestions are appreciated! Thank you!</p>