Derivative of covariance matrix

Mr_Cool: <p>Suppose I have a bivariate covariance function <code>c(x,y)</code> and at input <code>[x1,...,xn]</code> and <code>[y1,...,ym]</code> and I would like to compute the covariance matrix with <code>i,j</code>th element <code>d/dx c(xi,yj)</code>. Here I have assumed that each <code>xi</code> and <code>yj</code> are scalar valued, however this could easily be extended to finding the <code>k</code>th partial derivative of the covariance matrix. In general, I would like this to be extended to higher derivatives too (e.g. I would like to compute the matrix with <code>i,j</code>th element <code>d/dx d/dy c(xi,yj)</code>), but this may be an easy generalisation of the first issue.</p>
<p>I feel that PyTorch’s automatic differentiation system should be able to handle this, but I haven’t been able to come up with any nice solution. I think the issue is to do with broadcasting derivatives. For instance, to compute the non-differentiated covariance matrix, I would do something like <code>c(x.unsqueeze(1),y)</code>,  but the following code doesn’t work as I would like (using <code>autograd.grad</code>):</p>
<pre><code class="lang-auto">X = torch.Tensor([1,2,3]).requires_grad_(True).unsqueeze(1)
Y = torch.Tensor([4,5]).requires_grad_(True)
cov_matrix = c(X,Y)
dx_cov_matrix = grad(cov_matrix, [X], grad_outputs=[torch.ones([X.size(0),Y.size(0)])], create_graph=True)
</code></pre>
<p>Any help would be appreciated! Thanks</p>

pascal_notsawo: <p>Here is an implementation of the covariance with pytorch from <a href="https://discuss.pytorch.org/t/covariance-and-gradient-support/16217/2">covariance-and-gradient-support</a></p>
<pre><code class="lang-auto">import torch

def cov(m, rowvar=False, inplace=False):
    '''Estimate a covariance matrix given data.

    Thanks :
    - https://discuss.pytorch.org/t/covariance-and-gradient-support/16217/2
    - https://github.com/numpy/numpy/blob/master/numpy/lib/function_base.py#L2276-L2494

    Covariance indicates the level to which two variables vary together.
    If we examine N-dimensional samples, `X = [x_1, x_2, ... x_N]^T`,
    then the covariance matrix element `C_{ij}` is the covariance of
    `x_i` and `x_j`. The element `C_{ii}` is the variance of `x_i`.

    Args:
        m: A 1-D or 2-D array containing multiple variables and observations.
            Each row of `m` represents a variable, and each column a single
            observation of all those variables.
        rowvar: If `rowvar` is True, then each row represents a
            variable, with observations in the columns. Otherwise, the
            relationship is transposed: each column represents a variable,
            while the rows contain observations.

    Returns:
        The covariance matrix of the variables.
    '''
    if m.dim() &gt; 2:
        raise ValueError('m has more than 2 dimensions')
    if m.dim() &lt; 2:
        m = m.view(1, -1)
    if not rowvar and m.size(0) != 1:
        m = m.t()
    m = m.type(torch.double)  
    fact = 1.0 / (m.size(1) - 1)
    if inplace:
        m -= torch.mean(m, dim=1, keepdim=True)
    else :
        m = m - torch.mean(m, dim=1, keepdim=True)
    mt = m.t()  # if complex: mt = m.t().conj()
    return fact * m.matmul(mt).squeeze()
</code></pre>
<p>You can also calculate the gradient as follows.</p>
<pre><code class="lang-auto">def gradient(y, x, grad_outputs=None):
    """Compute dy/dx @ grad_outputs"""
    if grad_outputs is None:
        grad_outputs = torch.ones_like(y)
    grad = torch.autograd.grad(y, [x], grad_outputs=grad_outputs, create_graph=True)[0]
    return grad
</code></pre>
<p>Now let’s take an example.</p>
<pre><code class="lang-auto">x, y = [90, 60, 90], [90, 90, 30]
A = torch.tensor([x, y], dtype = torch.double, requires_grad=True)
c = cov(m = A, rowvar = True)
</code></pre>
<blockquote>
<p>tensor([[ 300., -300.], [-300., 1200.]], dtype=torch.float64, grad_fn=&lt;MulBackward0&gt;)</p>
</blockquote>
<p>You can now calculate the gradient of <code>c</code> w.r.t to <code>A</code>, like this.</p>
<pre><code class="lang-auto">gradient(y = c, x = A)
</code></pre>
<blockquote>
<p>tensor([[ 30.,   0., -30.], [ 30.,   0., -30.]], dtype=torch.float64, grad_fn=&lt;AddBackward0&gt;)</p>
</blockquote>
<p>But I guess this is not the behaviour you want, because <code>torch.autograd.grad</code> makes dimension reductions for multi-dimensional parameters.</p>
<p>So you can calculate <code>dc(xi, yj)/dx or dc(xi, yj)/dy</code> like this :</p>
<pre><code class="lang-auto">grad_cij = gradient(y = c[i][j], x = A)

dc(xi, yj)/dx = grad_cij[0]

dc(xi, yj)/dy = grad_cij[1]
</code></pre>
<pre><code class="lang-auto">i = 0
j = 0
dcij_dA = gradient(y = c[i][j], x = A)
</code></pre>
<blockquote>
<p>tensor([[ 10., -20.,  10.], [  0.,   0.,   0.]], dtype=torch.float64, grad_fn=)</p>
</blockquote>
<p>Since <code>c[0][0]</code> is independent of <code>y</code>, <code>dcij_dA[1]</code> is null. Similarly, you can see that the derivative of <code>c[1][1]</code> with respect to <code>x</code> is also null.</p>
<pre><code class="lang-auto">i = 1
j = 1
dcij_dA = gradient(y = c[i][j], x = A)
</code></pre>
<blockquote>
<p>tensor([[  0.,   0.,   0.], [ 20.,  20., -40.]], dtype=torch.float64, grad_fn=&lt;AddBackward0&gt;)</p>
</blockquote>
<p>You can also notice that <code>c[1][0] = c[0][1]</code>, all having non-zero derivatives with respect to <code>x</code> and <code>y</code>.</p>
<pre><code class="lang-auto">i = 0
j = 1
dcij_dA = gradient(y = c[i][j], x = A)
</code></pre>
<blockquote>
<p>tensor([[ 10.,  10., -20.], [  5., -10.,   5.]], dtype=torch.float64, grad_fn=&lt;AddBackward0&gt;)</p>
</blockquote>
<pre><code class="lang-auto">i = 1
j = 0
dcij_dA = gradient(y = c[i][j], x = A)
dcij_dA
</code></pre>
<blockquote>
<p>tensor([[ 10.,  10., -20.], [  5., -10.,   5.]], dtype=torch.float64, grad_fn=&lt;AddBackward0&gt;)</p>
</blockquote>
<p>You can also calculate everything like this :</p>
<pre><code class="lang-auto">n, m = c.shape
dc_dA = torch.stack([gradient(y = c[i][j], x = A) for j in range(m) for i in range(n)])
</code></pre>