Compute matrix divergence

yuanhao: <p>Is there any method (or what is the easiest way) to compute the matrix divergence.</p>
<p>The matrix divergence is given by the following (<a href="https://en.wikipedia.org/wiki/Divergence#Tensor_field" class="inline-onebox" rel="noopener nofollow ugc">Divergence - Wikipedia</a>, and a screenshot is also attached)</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/2/3/23bfaeed71032b8a51c0cfb39bbae81e9ddc2ccb.png" data-download-href="https://discuss.pytorch.org/uploads/default/23bfaeed71032b8a51c0cfb39bbae81e9ddc2ccb" title="Screenshot_20240915_111658"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/2/3/23bfaeed71032b8a51c0cfb39bbae81e9ddc2ccb_2_690x264.png" alt="Screenshot_20240915_111658" data-base62-sha1="56foffDAvS5iS7bWBo6r1BQsIn9" width="690" height="264" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/2/3/23bfaeed71032b8a51c0cfb39bbae81e9ddc2ccb_2_690x264.png, https://discuss.pytorch.org/uploads/default/optimized/3X/2/3/23bfaeed71032b8a51c0cfb39bbae81e9ddc2ccb_2_1035x396.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/2/3/23bfaeed71032b8a51c0cfb39bbae81e9ddc2ccb_2_1380x528.png 2x" data-dominant-color="F6F6F6"><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_20240915_111658</span><span class="informations">1519×583 96.7 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>

KFrank: <p>Hi Hao!</p>
<aside class="quote no-group" data-username="yuanhao" data-post="1" data-topic="209597" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/yuanhao/48/72757_2.png" class="avatar"> yuanhao:</div>
<blockquote>
<p>Is there any method (or what is the easiest way) to compute the matrix divergence.</p>
</blockquote>
</aside>
<p>Yes.  The short story is that you can use <code>func.jacfwd()</code> to<br>
compute the full jacobian of <code>A</code> with respect to <code>x</code> and then pick<br>
out the specific partial derivatives that you need for the matrix<br>
divergence.</p>
<p>First, a couple of comments:</p>
<p>Autograd performs differentiation with respect to entire tensors,<br>
not just specific elements of tensors.  So you can’t differentiate<br>
<code>A[1, 1]</code> with respect to <code>x[1]</code> without also differentiating it with<br>
respect to <code>x[2]</code>.  (Analogously, in forward-mode autograd, you<br>
can’t differentiate <code>A[1, 1]</code> with respect to <code>x[1]</code> without also<br>
differentiating<code> A[1, 2]</code> with respect to <code>x[1]</code>.)  So there is no<br>
practical way to compute just the specific individual partial<br>
derivatives that you need.</p>
<p>When you use backward-mode autograd (e.g., <code>.backward()</code>),<br>
autograd computes the gradient of a single scalar with respect<br>
to all of the elements of potentially multiple tensors.  When you<br>
use forward-mode autograd, it computes the directional<br>
derivatives of all of the elements of potentially multiple tensors<br>
with respect to a single direction within the tensor(s) with respect<br>
to which you are differentiating.</p>
<p>To compute a jacobian, you need the partial derivatives of all of<br>
the elements of your output tensor(s) with respect to all of the<br>
elements of your input tensor(s).  So, regardless of whether you<br>
are using backward- or forward-mode autograd, you will need<br>
multiple autograd passes, one for each output element you are<br>
differentiating, or one for each input element with respect to<br>
which you are differentiating.</p>
<p>In backward mode, you would need to perform nine autograd<br>
passes (one for each of the nine elements of <code>A</code>).  In forward<br>
mode you would need three autograd passes (one for each of<br>
the three elements of <code>x</code>).  So it is likely to be cheaper to use<br>
forward mode than backward mode for this use case.</p>
<p>You don’t have to run the multiple autograd passes by hand<br>
yourself – pytorch packages the full jacobian computation<br>
in <code>func.jacrev()</code> (backward mode) and <code>func.jacfwd()</code><br>
(forward mode).</p>
<p>Here is a script that computes your matrix divergence using<br>
<code>jacfwd()</code> (because that is likely to be more efficient):</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
print (torch.__version__)

def A_of_x (x):                        # some example function that mixes the elements of x together
    xsq = (x**2).sum()
    return  torch.outer (xsq * x, x)

jacA = torch.func.jacfwd (A_of_x)      # jacA is a functional that computes jacobian at a specific point

x = torch.arange (3.) + 1              # some point at which to compute jacobian
print (x)

jac = jacA (x)                         # computes full jacobian, not just desired elements
print (jac)

div = torch.einsum ('ijj -&gt; i', jac)   # use einsum to pick out the desired elements and sum
print (div)

# to see what's going on, compute div by picking out elements and summing by hand
div_mask = torch.eye (3).unsqueeze (1).expand (3, 3, 3)
print (div_mask)
print (div_mask * jac)
print ((div_mask * jac).sum ((0, 2)))
</code></pre>
<p>And here is its output:</p>
<pre><code class="lang-plaintext">2.4.1
tensor([1., 2., 3.])
tensor([[[ 30.,   4.,   6.],
         [ 32.,  22.,  12.],
         [ 48.,  12.,  32.]],

        [[ 32.,  22.,  12.],
         [  8.,  72.,  24.],
         [ 12.,  66.,  64.]],

        [[ 48.,  12.,  32.],
         [ 12.,  66.,  64.],
         [ 18.,  36., 138.]]])
tensor([ 84., 168., 252.])
tensor([[[1., 0., 0.],
         [1., 0., 0.],
         [1., 0., 0.]],

        [[0., 1., 0.],
         [0., 1., 0.],
         [0., 1., 0.]],

        [[0., 0., 1.],
         [0., 0., 1.],
         [0., 0., 1.]]])
tensor([[[ 30.,   0.,   0.],
         [ 32.,   0.,   0.],
         [ 48.,   0.,   0.]],

        [[  0.,  22.,   0.],
         [  0.,  72.,   0.],
         [  0.,  66.,   0.]],

        [[  0.,   0.,  32.],
         [  0.,   0.,  64.],
         [  0.,   0., 138.]]])
tensor([ 84., 168., 252.])
</code></pre>
<p>Best.</p>
<p>K. Frank</p>